Henry_Chibueze:

The training and weight updates only happened in the train function, so if u call the model any other place, it won’t have those trained weights and It’ll still be the untrained model it is…simply because you passed the model as a parameter to the train function and made the trained model peculiar to just the train function.


I don’t think that’s correct as seen e.g. here:
import torch
import torch.nn as nn

def train(model, optimizer):
    for epoch in range(10):
        optimizer.zero_grad()
        output = model(torch.randn(1, 1))
        loss = criterion(output, torch.randn(1, 1))
        loss.backward()
        optimizer.step()
        print('epoch {}, loss {}, weight {}'.format(
            epoch, loss.item(), model.weight))

if __name__=='__main__':
    model = nn.Linear(1, 1)
    print('before training')
    print(model.weight)
    
    optimizer = torch.optim.SGD(model.parameters(), lr=1.)
    criterion = nn.MSELoss()
    
    train(model, optimizer)
    print('after training')
    print(model.weight)

As you can see, the model is created in the if-clause guard, passed to the train function, and is not returned.
Nevertheless, the weight updates will be performed inplace on the model reference and thus the model in the same scope will contain the upgraded parameters.