Here’s my classifier -
class Classifier(nn.Module):
    def __init__(self, num_classes):
        super(Classifier(, self).__init__()
        
        
        self.base_conv = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2)
        
        self.conv_block_1 = ConvBlock(in_channels=16, out_channels=32)
        self.conv_block_2 = ConvBlock(in_channels=32, out_channels=64)
                
        self.lastcnn = nn.Conv2d(in_channels=64, out_channels=num_classes, kernel_size=28, stride=1, padding=0)
        
        self.gradients = None
        
        
    def forward(self, x):
        
        x = self.base_conv(x)
        
        x = self.conv_block_1(x)
        x = self.conv_block_2(x)
                                       
        if x.requires_grad: x.register_hook(self.store_gradients)

        x = self.lastcnn(x)
        
        return x
      
      
    def store_gradients(self, grad):
        self.gradients = grad    
    
    def get_gradients(self):
        return self.gradients

After I training this model, I save it as a dict using torch.save(model.state_dict(), PATH).
On loading the model, when I try to access it’s get_gradient method as follows, no gradients are returned.
model.eval()
pred = model(single_image.to(device))
pred = torch.softmax(pred, dim = 1)
pred_idx = pred.argmax(dim=1).item()

gradients = model.get_gradients()
gradients  #this is empty

On the flipside, If run the exact same code after training (instead of saving and loading), gradients are returned.
I tried saving the whole model using torch.save() and then load it again. The code works. So, it’s because of the parameters not getting saved in the model dict. Any way to store all the additional parameters when doing torch.save(model.state_dict(), PATH)?