I build my custom layer which acts like a rnn, but is has more states than a regular rnn cell. I’ke like to record the graph using tensorboard. The training is successful, and i don’t see error in my code, however, when i call writer.add_graph(mymodel,dataloader), it throws an error
File "/home/mypc/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 531, in _slow_forward
    result = self.forward(*input, **kwargs)
TypeError: forward() missing 1 required positional argument: 'state_6'"

Here is a minimal example to produce this error. This model just executes arbitrary operation, so it looks doesn’t make sense at all, please ignore this. But it reproduces the error.
import torch
import numpy as np
import random
import matplotlib.pyplot as plt
import utilities
import importlib
import matplotlib
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter

matplotlib.use('Qt5Agg')
# from tensorboardX import SummaryWriter


# parameters
input_size = 40
layer1_size = 100
layer2_size = 10
length = 200
batch_size = 16
epoch = 10

class RandPatternDataset(Dataset):
    """random pattern dataset"""

    def __init__(self,transform=None):
        self.dataset = np.random.rand(100,40,200).astype(np.float32)
        self.label = np.random.rand(100,10).astype(np.float32)
        self.transform = transform

    def __len__(self):
        return self.dataset.shape[0]

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        return self.dataset[idx],self.label[idx]

class custom_layer(torch.nn.Module):
    def __init__(self, input_size, layer_size, length):
        super().__init__()
        self.input_size = input_size
        self.layer_size = layer_size
        self.length = length
        self.batch_size = batch_size
        self.weight = torch.nn.Linear(input_size, layer_size, bias=True)

    def forward(self, inp, state_1, state_2, state_3, state_4, state_5, state_6):

        state_1 = state_1 * 0.1 + inp
        state_2 = state_2 * 0.1 + inp
        state_3 = state_1 + state_2
        state_4 = state_5*0.1
        state_5 = self.weight(state_3)
        state_6 = state_5*0.1
        out = state_5*0.1

        return out, state_1, state_2, state_3, state_4, state_5, state_6

rand_pattern = RandPatternDataset()
dataloader = DataLoader(rand_pattern, batch_size=batch_size, shuffle=True, drop_last=True)

torch.manual_seed(2)

spike_before_train = None
v_before_train = None

# define two layers
layer_1 = custom_layer(input_size, layer1_size, length)
layer_2 = custom_layer(layer1_size, layer2_size, length)

writer = SummaryWriter(log_dir="~/logdir")
for j in range(epoch):

    print(j)
    for i_batch, sample_batched in enumerate(dataloader):

        l1_s1 = torch.zeros([batch_size, input_size])
        l1_s2 = torch.zeros([batch_size, input_size])
        l1_s3 = torch.zeros([batch_size, input_size])
        l1_s4 = torch.zeros([batch_size, layer1_size])
        l1_s5 = torch.zeros([batch_size, layer1_size])
        l1_s6 = torch.zeros([batch_size, layer1_size])

        l2_s1 = torch.zeros([batch_size, layer1_size])
        l2_s2 = torch.zeros([batch_size, layer1_size])
        l2_s3 = torch.zeros([batch_size, layer1_size])
        l2_s4 = torch.zeros([batch_size, layer2_size])
        l2_s5 = torch.zeros([batch_size, layer2_size])
        l2_s6 = torch.zeros([batch_size, layer2_size])

        # loss
        criterion = torch.nn.MSELoss()

        # optimizer

        params = list(layer_1.parameters()) + list(layer_2.parameters())
        optimizer = torch.torch.optim.Adam(params, lr=0.0005, betas=(0.9, 0.999), eps=1e-08)

        x_train = sample_batched[0]
        target = sample_batched[1]

        # step computation
        for t in range(length):
            out1, l1_s1,l1_s2,l1_s3,l1_s4,l1_s5,l1_s6 = layer_1(x_train[:, :, t], l1_s1,l1_s2,l1_s3,l1_s4,l1_s5,l1_s6)

            out2, l2_s1,l2_s2,l2_s3,l2_s4,l2_s5,l2_s6 = layer_2(out1,l2_s1,l2_s2,l2_s3,l2_s4,l2_s5,l2_s6)

        layer_1.zero_grad()
        layer_2.zero_grad()

        loss = criterion(out2, target)

        loss.backward()

        print(loss)

        optimizer.step()

writer.add_graph(layer_2,dataloader)
writer.close()

I also tried another implementation. I created a new module, which is custom_model, and put layer_1, layer_2 and the computation inside the forward() of custom_model. it does the same thing as above code. it causes a different error:
  File "/home/mypc/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 531, in _slow_forward
    result = self.forward(*input, **kwargs)
TypeError: forward() takes 2 positional arguments but 7 were given

here is the implementation of custom_model
class custom_model(torch.nn.Module):
    def __init__(self):

        super(custom_model, self).__init__()

        self.length = length
        self.layer_1 = custom_layer(input_size, layer1_size, length)
        self.layer_2 = custom_layer(layer1_size, layer2_size, length)

    def forward(self, inp):
	
        l1_s1 = torch.zeros([batch_size, input_size])
        l1_s2 = torch.zeros([batch_size, input_size])
        l1_s3 = torch.zeros([batch_size, input_size])
        l1_s4 = torch.zeros([batch_size, layer1_size])
        l1_s5 = torch.zeros([batch_size, layer1_size])
        l1_s6 = torch.zeros([batch_size, layer1_size])

        l2_s1 = torch.zeros([batch_size, layer1_size])
        l2_s2 = torch.zeros([batch_size, layer1_size])
        l2_s3 = torch.zeros([batch_size, layer1_size])
        l2_s4 = torch.zeros([batch_size, layer2_size])
        l2_s5 = torch.zeros([batch_size, layer2_size])
        l2_s6 = torch.zeros([batch_size, layer2_size])

        for t in range(length):
            out1, l1_s1,l1_s2,l1_s3,l1_s4,l1_s5,l1_s6 = self.layer_1(inp[:, :, t], l1_s1,l1_s2,l1_s3,l1_s4,l1_s5,l1_s6)
            out2, l2_s1,l2_s2,l2_s3,l2_s4,l2_s5,l2_s6 = self.layer_2(out1,l2_s1,l2_s2,l2_s3,l2_s4,l2_s5,l2_s6)

        return  out2

mymodel = custom_model()
writer = SummaryWriter(log_dir="~/logdir")
for j in range(epoch):

    print(j)
    for i_batch, sample_batched in enumerate(dataloader):
        # loss
        criterion = torch.nn.MSELoss()
        # optimizer
        params = list(mymodel.parameters())
        optimizer = torch.torch.optim.Adam(params, lr=0.0005, betas=(0.9, 0.999), eps=1e-08)

        x_train = sample_batched[0]
        target = sample_batched[1]

        filtered_output = mymodel.forward(x_train)

        mymodel.zero_grad()

        loss = criterion(out2, target)

        loss.backward()
        print(loss)
        optimizer.step()

writer.add_graph(mymodel,dataloader)
writer.close()

I’m using pytorch 1.2.0, python 3.6.9, tensorboard 1.14.0