TypeError: forward() missing 1 required positional argument with Tensorboard PyTorch
I am trying to write my model to tensorboard with the following code:

model = SimpleLSTM(4, HIDDEN_DIM, HIDDEN_LAYERS, 1, BATCH_SIZE, device)
writer = tb.SummaryWriter(log_dir=tb_path)
sample_data = iter(trainloader).next()[0]
writer.add_graph(model, sample_data.to(device))


I get the error: TypeError: forward() missing 1 required positional argument: 'batch_size'

My model looks like this:

class SimpleLSTM(nn.Module):

    def __init__(self, input_dims, hidden_units, hidden_layers, out, batch_size, device):
        super(SimpleLSTM, self).__init__()
        self.input_dims = input_dims
        self.hidden_units = hidden_units
        self.hidden_layers = hidden_layers
        self.batch_size = batch_size
        self.device = device
        self.lstm = nn.LSTM(self.input_dims, self.hidden_units, self.hidden_layers,
                            batch_first=True, bidirectional=False)
        self.output_layer = nn.Linear(self.hidden_units, out)

    def init_hidden(self, batch_size):

        hidden = torch.rand(self.hidden_layers, batch_size, self.hidden_units, device=self.device, dtype=torch.float32)
        cell = torch.rand(self.hidden_layers, batch_size, self.hidden_units, device=self.device, dtype=torch.float32)
        hidden = nn.init.xavier_normal_(hidden)
        cell = nn.init.xavier_normal_(cell)
        return (hidden, cell)

    def forward(self, input, batch_size):
        hidden = self.init_hidden(batch_size)  incomplete batch
        lstm_out, (h_n, c_n) = self.lstm(input, hidden)
        raw_out = self.output_layer(h_n[-1])
        return raw_out


How can I write this model to TensorBoard?

Your model takes two arguments input and batch_size, but you only provide one argument for add_graph to call your model with.

The inputs (second argument to add_graph) should be a tuple with the input and the batch_size:

writer.add_graph(model, (sample_data.to(device), BATCH_SIZE))


You don't really need to provide the batch size to the forward method, because you can infer it from the input. As your LSTM uses batch_first=True, it means that the input is required to have size [batch_size, seq_len, num_features], therefore the size of the first dimension is the current batch size.

def forward(self, input):
    batch_size = input.size(0)
    # ...

