Update comfy/float.py

comfy/float.py

@@ -60,8 +60,9 @@ def stochastic_rounding(value, dtype, seed=0):
         output = torch.empty_like(value, dtype=dtype)
         num_slices = max(1, (value.numel() / (4096 * 4096)))
         slice_size = max(1, round(value.shape[0] / num_slices))
-        for i in range(0, value.shape[0], slice_size):
-            output[i:i+slice_size].copy_(manual_stochastic_round_to_float8(value[i:i+slice_size], dtype, generator=generator))
+        with torch.no_grad():
+            for i in range(0, value.shape[0], slice_size):
+                output[i:i+slice_size].copy_(manual_stochastic_round_to_float8(value[i:i+slice_size], dtype, generator=generator))
         return output
 
     return value.to(dtype=dtype)