I’m using the double backward capability for BatchNorm that was recently added (thanks so much!). But I’m getting an error that I’m struggling to track down. I’ve tried to write down a minimal example.
import torch import torch.nn as nn from torch.autograd import Variable class Net(nn.Module): def __init__(self): super(self.__class__, self).__init__() self.classifier = nn.Sequential( nn.Linear(4,4), nn.BatchNorm1d(4, momentum=0.1, affine=False), nn.ReLU(inplace=True) ) def forward(self, x): return self.classifier(x) input_var = Variable(torch.ones((4))).cuda() net = Net() net.cuda() out = net.forward(input_var) loss = sum(out) loss.backward(create_graph=True) loss.backward(create_graph=True)
The output of running this code is an error:
RuntimeError: missing required argument 'save_mean'
For some reason, I’m not able to trace the error any further than the
torch.autograd.backward() function. And the only place
save_mean seems to appear is in the legacy batch norm code.
Can anyone see what I’m doing wrong?