Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- # Train the model
- total_step = len(train_loader)
- for epoch in range(num_epochs):
- for i, (images, labels) in enumerate(train_loader):
- # Move tensors to the configured device
- images = images.to(device)
- labels = labels.to(device)
- # Forward pass
- outputs = model(images)
- loss = loss_function(outputs, labels)
- # Backward and optimize
- optimizer.zero_grad()
- loss.backward()
- optimizer.step()
- if (i+1) % 100 == 0:
- print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
- .format(epoch+1, num_epochs, i+1, total_step, loss.item()))
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement