Skip to content

Instantly share code, notes, and snippets.

@Rishav09
Created August 8, 2022 21:56
Show Gist options
  • Save Rishav09/f52dd02776bb6bf9fed792b526de9e41 to your computer and use it in GitHub Desktop.
Save Rishav09/f52dd02776bb6bf9fed792b526de9e41 to your computer and use it in GitHub Desktop.
model.train()
for batch_idx, (data, target) in enumerate(loader['train']):
# move to GPU
if torch.cuda.is_available():
data, target = data.to('cuda', non_blocking=True), target.to('cuda', non_blocking = True) # noqa
optimizer.zero_grad()
output = model(data)
pred = torch.argmax(output, dim=1)
# pred = pred.to(torch.double)
# target = target.to(torch.double)
loss = criterion(pred.float(), target.float())
loss.backward()
optimizer.step()
train_loss += ((1 / (batch_idx + 1)) * ((loss.data) - train_loss))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment