@ColdFrenzy
nn.Sequential should work as mentioned by @thecho7. Below is an example
ORIGINAL
import numpy as np
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable

class random_model(nn.Module):
    def __init__(self, num_layers):
        super(random_model, self).__init__()
        self.model1 = nn.Sequential(nn.Linear(100, 20), nn.BatchNorm1d(20), nn.ReLU())
        self.model2 = []
        for i in range(num_layers):
            self.model2.append(nn.Sequential(nn.Linear(20, 20), nn.BatchNorm1d(20), nn.ReLU()))
        #self.model2 = nn.Sequential(*self.model2)
        self.model3 = nn.Linear(20, 1)
    
    def forward(self, x):
        x = self.model1(x)
        for cur_model in self.model2:
            x = cur_model(x)
        #x = self.model2(x)
        x = self.model3(x)
        return x
    
model = random_model(10)
loss = torch.nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

X = torch.rand(100, 100)
y = torch.rand(100)

import wandb
wandb.init(project="random-project-custom-class", entity="")
wandb.config = {
  "learning_rate": 0.001,
  "epochs": 100,
  "batch_size": 100
}

# 3. Log gradients and model parameters
wandb.watch(model, log="gradients", log_freq=100)
print("DEF")
for cur_epoch in range(100):
    optimizer.zero_grad()
    output = model(X)
    cur_loss = loss(output, y)
    cur_loss.backward()
    optimizer.step()
    cur_loss = loss(output, y)
    print("Epoch {0} Loss is {1}".format(cur_epoch, cur_loss))
    wandb.log({"loss": cur_loss})


image1166×595 28.7 KB

MODIFIED
import numpy as np
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable

class random_model(nn.Module):
    def __init__(self, num_layers):
        super(random_model, self).__init__()
        self.model1 = nn.Sequential(nn.Linear(100, 20), nn.BatchNorm1d(20), nn.ReLU())
        self.model2 = []
        for i in range(num_layers):
            self.model2.append(nn.Sequential(nn.Linear(20, 20), nn.BatchNorm1d(20), nn.ReLU()))
        self.model2 = nn.Sequential(*self.model2)
        self.model3 = nn.Linear(20, 1)
    
    def forward(self, x):
        x = self.model1(x)
        #for cur_model in self.model2:
        #    x = cur_model(x)
        x = self.model2(x)
        x = self.model3(x)
        return x
    
model = random_model(10)
loss = torch.nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

X = torch.rand(100, 100)
y = torch.rand(100)

import wandb
wandb.init(project="random-project-custom-class", entity="")
wandb.config = {
  "learning_rate": 0.001,
  "epochs": 100,
  "batch_size": 100
}

# 3. Log gradients and model parameters
wandb.watch(model, log="gradients", log_freq=100)
print("DEF")
for cur_epoch in range(100):
    optimizer.zero_grad()
    output = model(X)
    cur_loss = loss(output, y)
    cur_loss.backward()
    optimizer.step()
    cur_loss = loss(output, y)
    print("Epoch {0} Loss is {1}".format(cur_epoch, cur_loss))
    wandb.log({"loss": cur_loss})


image1301×744 61.5 KB