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`

.