Visualizing LSTM Model with Tensorboard

I am trying to visualize my LSTM model, which looks like this:

class SequenceModel(nn.Module):

  def __init__(self, n_features, n_classes, n_hidden=100, n_layers=3):

    self.lstm = nn.LSTM( 
        input_size = n_features,
        hidden_size = n_hidden,
        num_layers = n_layers,
        batch_first = True,
        dropout = 0.945
    self.classifier = nn.Linear(n_hidden, n_classes)

  def forward(self,x):
    _, (hidden,_) = self.lstm(x)

    out = hidden [-1] 
    return self.classifier(out) 

 (model): SequenceModel(
    (lstm): LSTM(3, 100, num_layers=3, batch_first=True, dropout=0.945)
    (classifier): Linear(in_features=100, out_features=2, bias=True)
  (criterion): CrossEntropyLoss()

I know usually you do this like this:

To write the computational graph we will be using add_graph() method. add_graph requires two arguments

  1. The model
  2. A sample image for the same shape as that of the input to track how it changes as it passes through the network

So i tried it with this code:

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()
writer.add_graph(model,  sample)

But I am getting this error:

Only tensors, lists, tuples of tensors, or dictionary of tensors can be output from traced functions
Error occurs, No graph saved

I canĀ“t figure out what the problem is, has anyone an idea? Thanks in advance.