My model is:
import torch.nn as nn

class Net(nn.Module):
    def __init__(self, num_input, num_hidden, num_classes, dropout,
                 activation='tanh'):
        super(Net, self).__init__()
        self.dropout = nn.Dropout(dropout)
        self.fc1 = nn.Linear(num_input, num_hidden)
        self.fc2 = nn.Linear(num_hidden, num_classes)

        if activation == 'tanh':
            self.activation_f = torch.tanh
        elif activation == 'relu':
            self.activation_f = torch.relu

    def forward(self, x):
        x = self.activation_f(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

I call my model for instance as:
model = Net(14,512,2,0.2).to(device)

However once I use TorchScript as:
traced_model = torch.jit.trace(model, torch.zeros([1, 14], dtype=torch.float))

I receive the following error:
IndexError: The shape of the mask [2] at index 0 does not match the shape of the indexed tensor [1, 2] at index 0

I know that if I use model.eval() I don’t receive any error BUT I want to use my model for training and not evaluation. Does anybody know any solution or workaround for such problem?
PS: I am using PyTorch version 1.4.