winglian
/

mistral-denseformer-7b

Text Generation

mistral_denseformer

Model card Files Files and versions Community

winglian commited on Mar 24

Commit

4d34a45

•

1 Parent(s): e62a1d0

Update denseformer.py

Files changed (1) hide show

denseformer.py +1 -1

denseformer.py CHANGED Viewed

@@ -7,7 +7,7 @@ class InPlaceSetSlice(torch.autograd.Function):
   def forward(ctx, full_tensor, last_slice, x_idx, x_val):
     full_tensor[x_idx] = x_val
     ctx.x_idx = x_idx
-    ret = torch.Tensor().to(full_tensor.device)
     ret.set_(full_tensor[:x_idx + 1])
     return ret

   def forward(ctx, full_tensor, last_slice, x_idx, x_val):
     full_tensor[x_idx] = x_val
     ctx.x_idx = x_idx
+    ret = torch.Tensor().to(device=full_tensor.device, dtype=torch.bfloat16)
     ret.set_(full_tensor[:x_idx + 1])
     return ret