Ok cool, what about if it’s set up this way.
crit = nn.SomeLoss()
optim = optim.SGD()
net = models.resnet18()

for e in range(num_epochs):
    
    # training
    pred = net(some_data)
    optim.zero_grad()
    loss = crit(pred, target)
    loss.backward() 
    optim.step()

    # validation
    valid_pred = net(some_validation_data)
    loss = crit(valid_pred, valid_target)

Would zero_grad take care of that?