Hello, I’m trying to implement a model that is used multiple times in one step of each epoch. I think the code is like below. How to train such model in proper way? Thank you.
encoder = Encoder()
optimizer = torch.optim.SGD(encoder.parameters())
criterion = torch.nn.MSELoss()
for epoch in range(num_epoch):
    for images, labels in train_loader:
        # in this line, images.shape [b, c, h, w]
        something processing ...
        # in this line, images.shape [n, b, c, h, w]
        outs = []
        for _ in images[:]:
            outs = encoder(_)
        something processing ...
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()