@apaszke Does the following look like a correct implementation?
If yes, do you think JIT / TorchScript could bring more performance improvements on CUDA (e.g. fuse this with Conv2d or BatchNorm1d)? I saw @ngimel used torch.rand instead of torch.bernoulli for better fusability, but does this apply to a custom autograd function?
class ReLUDropoutInplace(torch.nn.Module):
    def __init__(self, p):
        super(ReLUDropoutInplace, self).__init__()
        self.p = p

    def forward(self, input):
        if self.training:
            #p1m = 1. - self.p
            #mask = torch.rand_like(input) < p1m
            #mask *= (input > 0)
            #return mask.type_as(input) * input * (1./p1m)
            return ReLUDropoutInplace.ReLUDropoutInplaceFunction.apply(input, self.p)
        else:
            return input.clamp_(min = 0)

    class ReLUDropoutInplaceFunction(torch.autograd.Function):
        @staticmethod
        def forward(ctx, input, p):
            ctx.p1m = 1. - p
            mask = torch.rand_like(input) < ctx.p1m
            mask *= (input > 0)
            one_minus_mask = 1 - mask
            ctx.save_for_backward(one_minus_mask)

            input = input.contiguous().masked_fill_(one_minus_mask, 0)
            input *= (1. / ctx.p1m)

            return input

        @staticmethod
        def backward(ctx, grad):
            one_minus_mask, = ctx.saved_tensors
            grad = grad.contiguous().masked_fill_(one_minus_mask, 0)
            grad *= 1. / ctx.p1m
            return grad, None

If I try the following, I’m getting RuntimeError: a leaf Variable that requires grad has been used in an in-place operation. at the masked_fill_ operation:
if self.training:
            p1m = 1. - self.p
            mask = torch.rand_like(input) < p1m
            mask *= (input > 0)
            one_minus_mask = 1 - mask
            input.masked_fill_(one_minus_mask, 0)
            input *= (1. / ctx.p1m)
            return input

...
if __name__ == '__main__':
    x = torch.rand(2, 16, 64, 128)
    x = (x + 1).requires_grad_()
    m = ReLUDropoutInplace(p = 0.2)
    m.train()

    y = m(x)
    y.sum().backward()