Update llama.py
This commit is contained in:
parent
8123210d7e
commit
6364232757
1 changed files with 1 additions and 1 deletions
|
|
@ -937,7 +937,7 @@ def LlamaModel_fast_forward_inference(
|
|||
next_decoder_cache = []
|
||||
residual = torch.empty_like(X)
|
||||
print(bsz, q_len, hd)
|
||||
_XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32)
|
||||
_XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
|
||||
XX, XX2 = _XX[0], _XX[1]
|
||||
variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = "cuda:0")
|
||||
temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue