Hello, I am using torch.jit.script to export my models but I encountered a problem with Resnet. This is my model definition:
class Model(nn.Module):
    def __init__(self, num_cats):
        super(Model, self).__init__()    
        self.model = torchvision.models.resnet18(pretrained=True)
        self.model.fc = nn.Linear(self.model.fc.in_features, num_cats)
        self.model = torch.jit.trace(self.model, torch.rand(1,3,224,224))
    def forward(self, x):
        x = self.model(x)
        return x

Training works fine and validation accuracy goes as expected. Finally I export my model like this
scripted = torch.jit.script(model)
torch.jit.save(scripted, 'scripted_model.pth')

The thing is that when I load my model
model = torch.jit.load('scripted_model.pth')

and use it to perform inference on the entire validation set with the same batch size than the one used for training I get the same predictions. But if I try to perform inference on a single image (or in the validation set with a small batch size), I get totally wrong predictions.
I’ve tried the same example with a custom CNN and I get good results. Also I tried exporting the resnet model with torch.jit.trace and torch.jit.trace_module (instead of torch.jit.script) and also get good results.
I would appreciate any help on the issue, since I find the functionality of torch.jit.script very powerful in production.