Hello, I am new to the community. Excited to be here.
As i am a new user of pytorch facing a similar problem, could someone provide an example on how the train method would look like inside the class?
For example in my code below i call self.forward() in the fit method, but is that correct? Calling self.model throws the following error: TypeError: forward() takes 1 positional argument but 2 were given
class AE(nn.Module):

    def __init__(self,h=[500,300,100,300,500], x_dim=None, dropout=0.1, lr=0.1):
        super(AE,self).__init__()

        self.x_dim = x_dim
        self.h = h
        self.n_layers = len(h)
        self.dropout = dropout
        self.lr = lr

        x=self.x_dim
        self.model = nn.ModuleDict()
        for i in np.arange(self.n_layers):
            n_nodes = self.h[i]
            if i == int(self.n_layers/2):

                self.model['latent_layer'] = nn.Linear(x, n_nodes, bias=True)
            else:
                self.model['hidden%d'%i] = nn.Linear(x, n_nodes, bias=True)
            x = n_nodes

        self.output = nn.Linear(x, self.x_dim,bias=True)

    def forward(self,x):

        for name,layer in self.model.items():
            x = F.relu(layer(x))

        output = self.output(x) # linear output
        return output


    def optimizer(self):

        return optim.Adam(self.model.parameters(),lr=self.lr)

   def fit(self,loader, epochs=10, batch_size=128):

        loss = RMSELoss()
        opt = self.optimizer()
        n_minibatches = len(loader)

        print('training on %d samples'%(n_minibatches*batch_size))
        for epoch in range(epochs):
            epoch_loss = 0

            for batch_x, batch_y in loader:
                #print(batch_x.shape, batch_y.shape)
                prediction = self.forward(batch_x.float()) 
                batch_loss = loss(prediction,batch_y.float())

                opt.zero_grad()   # clear gradients for next train
                batch_loss.backward()         # backpropagation, compute gradients
                opt.step()        # apply gradients

                epoch_loss += batch_loss.item()

            print('[Epoch %d] loss: %.3f'%(epoch+1, epoch_loss/n_minibatches))

Any help would be appreciated!
Thank you for your time