I have a custom module, which I am adding after VGG-16 feature module
class Last_layers_mod(nn.Module): def __init__(self): super(Last_layers_mod, self).__init__() self.gap = nn.AdaptiveAvgPool2d((1,1)) def forward(self, x): if self.train: x.register_hook(lambda x: print("gap reg_hook: ", x.size())) x = self.gap(x) x = x.view(-1, self.num_flat_features(x)) return x def num_flat_features(self, x): size = x.size()[1:] # all dimensions except the batch dimension num_features = 1 for s in size: num_features *= s return num_features
The model works perfectly during training but as soon as the first epoch ends and the model is in model.eval() phase, it throws
RuntimeError: cannot register a hook on a tensor that doesn’t require gradient
How to fix this issue?