Don't have optimizer.zero_grad() in your tutorial
hz1z opened this issue · 0 comments
hz1z commented
in your colab tutorial don;t have optimizer.zero_grad()
def train_epoch(model, loader, optimizer, epoch, loss_func):
model.train()
start_time = time.time()
run_loss = AverageMeter()
for idx, batch_data in enumerate(loader):
data, target = batch_data["image"].to(device), batch_data["label"].to(device)
logits = model(data)
loss = loss_func(logits, target)
loss.backward()
optimizer.step()
run_loss.update(loss.item(), n=batch_size)
print(
"Epoch {}/{} {}/{}".format(epoch, max_epochs, idx, len(loader)),
"loss: {:.4f}".format(run_loss.avg),
"time {:.2f}s".format(time.time() - start_time),
)
start_time = time.time()
return run_loss.avg