Understanding model.to(device)

Hello,

I am new to PyTorch, I just want to ensure that I correctly understand how model.to(device=device) works before I start creating more complex models.

So far, almost all the tutorials that Ive followed create, train and evaluate their models in the same notebook/script.

However, what if I am creating my model, training it, and evaluating it using functions/operations in different scripts?

Suppose I have the following code under these different scripts.

main.py

model = EfficientNet.from_pretrained("efficientnet-b3")
model = model.to(device=device)

for epoch in range(EPOCHS):
	train(model, train_loader, ...)
	check_accuracy(model, val_loader, ...)

train_model.py

def train_model(model, loader, ...):
	losses = []	
	for batch_idx, (data, targets) in enumerate(tqdm(loader)):
		...		
		scores=model(data)		
		...
	print(f"Loss: {sum(losses)/len(losses)}")

evaluation.py

def check_accuracy(model, loader, ...):
	model.eval()
	with torch.no_grad():
		...
		...
	model.train()
	return accuracy

Since my model is moved to device by running .to(device=device), I dont need to return the model from each function in different modules to main.py, is that right?

Simply passing the model to train_model() and check_accuracy() is enough run operations on the model, and once they finish running, the remaining operations in main.py will resume working on the same model, is that right?

Thank you all for your help.

Yes, your assumption should be correct as also seen in this post, since the model reference would be passed and its parameters (and buffers) updated inplace.
You could use the code snippet in the linked post to verify it using your setup.

1 Like

Thank you @ptrblck, will definitely check out the post.