Thanks for the detailed solution; however, I’ve just realized that my project is slightly different. My two kinds of training process. (1) training on the new layers. (2) training on old layers. They are both trained on different batches with some slightly different train procedures. One requires a loop in one dataset; another doesn’t
They would be like something like this.


# Train on new layers requires a loop on a dataset
for data in dataset_1():
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

 # Train on all layers doesn't loop the dataset
optimizer.zero_grad()
output = model(dataset2)
loss = criterion(output, target)
loss.backward()
optimizer.step()