dropout fix backward return + move into functional

memsave_torch/nn/Dropout.py

@@ -6,6 +6,8 @@
 import torch
 import torch.nn as nn
 
+from memsave_torch.nn.functional import dropoutMemSave
+
 
 class MemSaveDropout(nn.Dropout):
     """MemSaveDropout."""
@@ -42,38 +44,3 @@ def from_nn_dropout(cls, dropout: nn.Dropout):
         obj = cls(dropout.p)
         return obj
 
-
-# TODO: inplace
-class _MemSaveDropout(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x, p, train):
-        out, mask = torch.ops.aten.native_dropout(x, p, train)
-        if ctx.needs_input_grad[0]:
-            ctx.p = p
-            ctx.mask = mask
-        return out
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        grad_x = None
-
-        if ctx.needs_input_grad[0]:
-            grad_x = torch.ops.aten.native_dropout_backward(
-                grad_output, ctx.mask, scale=1 / (1 - ctx.p)
-            )
-
-        return grad_x
-
-
-def dropoutMemSave(x, p, training):
-    """Functional form of the memory saving dropout.
-
-    Args:
-        x: Input to the network
-        p: Probability of elements being zeroed
-        training: Whether the layer is in training mode (no dropout applied in eval)
-
-    Returns:
-        torch.Tensor: Output of the network
-    """
-    return _MemSaveDropout.apply(x, p, training)

memsave_torch/nn/functional/Dropout.py

@@ -25,7 +25,7 @@ def backward(ctx, grad_output):
                 grad_output, ctx.mask, scale=1 / (1 - ctx.p)
             )
 
-        return grad_x
+        return grad_x, None, None
 
 
 def dropoutMemSave(x, p, training) -> torch.Tensor: