Update denseformer.py

denseformer.py

@@ -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 = torch.Tensor().to(device=full_tensor.device, dtype=torch.bfloat16)
     ret.set_(full_tensor[:x_idx + 1])
     return ret