Register_hook throwing an error during model.eval()

autograd
1

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?

1 Like
2

You can replace if self.train: by if self.train and x.requires_grad: