dropout improvements

(cherry picked from commit ea5e6cb)
plutonium-239 · Aug 3, 2024 · 9d293fe · 9d293fe
1 parent 0b3b8c1
commit 9d293fe
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/memsave_torch/nn/functional/Dropout.py b/memsave_torch/nn/functional/Dropout.py
@@ -29,9 +29,10 @@ def backward(ctx, grad_output):
             mask = torch.empty_like(grad_output)
             mask = mask.bernoulli_(0.5).bool()
             torch.set_rng_state(orig_rng)
-            grad_x = torch.ops.aten.native_dropout_backward(
-                grad_output, mask, scale=1 / (1 - ctx.p)
-            )
+            grad_x = grad_output*mask/(1-ctx.p)
+            # grad_x = torch.ops.aten.native_dropout_backward(
+            #     grad_output, mask, scale=1 / (1 - ctx.p)
+            # )
 
         return grad_x, None, None