Hi, I am using my own custom std function for some reason.
I used to let keepdim=False as default, and everything worked just fine. However, I find that if keepdim=True, the gpu memory usage just keep going up every iteration and finally explode.
The function looks like this:
class CustomStd(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, dim=None, keepdim=False, eps=1e-5, unbiased=True):
        if dim is None:
            dim = tuple([i for i in range(input.dim())])
        dev = input - input.mean(dim=dim, keepdim=True)
        ctx.save_for_backward(input)
        ctx.eps=eps
        ctx.dev = dev
        ctx.numdim = input.dim()
        ctx.numel = functools.reduce(lambda x, y: x * y, [input.size(d) for d in dim])
        if unbiased:
            ctx.numel -= 1
        ctx.std = torch.sqrt(torch.sum(dev * dev, dim=dim, keepdim=True) / ctx.numel)
        return ctx.std if keepdim else ctx.std.squeeze()

    @staticmethod
    def backward(ctx, grad_output):
        input,= ctx.saved_tensors
        grad_input = grad_output
        for i in range(grad_output.dim(), ctx.numdim):
            grad_input = grad_input.unsqueeze(i)
        grad_input = ctx.dev * (ctx.numel - 1) / (ctx.numel**2) / (ctx.std + ctx.eps) * grad_input
        return grad_input, None, None, None, None

I assume that torch.squeeze() creates views onto the tensor and therefore the memory usage should be the same.
Am I missing anything here?
Thanks in advance.