What is the best and right way to do forward and backprop during model training and evaluation

distributed
1

1 ) model.train()
output = model(cat_feat_data, num_feat_data)
loss = criterion( output, target ) # Reshape for regression
epoch_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()

  1. model.train()
    optimizer.zero_grad()
    output = model(cat_feat_data, num_feat_data)
    loss = criterion( output, target ) # Reshape for regression
    epoch_loss += loss.item()
    loss.backward()
    optimizer.step()

  2. model.train()
    output = model(cat_feat_data, num_feat_data)
    loss = criterion( output, target ) # Reshape for regression
    epoch_loss += loss.item()
    loss.backward()
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

2

@Sajin_P

The best way to do would be like this
model.train()
optimizer.zero_grad()
output = model(cat_feat_data, num_feat_data)
loss = criterion( output, target ) # Reshape for regression
loss.backward()
optimizer.step()

epoch_loss += loss.item()

During evaluation , we dont do back propagation , we only do forward propagation

1 Like