<jupyter_start><jupyter_text>Train<jupyter_code>device = 'gpu' if torch.cuda.is_available() else 'cpu'
device
def train(epochs, train_dl, val_dl, model, criterion, hooks, lr=3e-3, 
          metrics=None, verbose=True):
    for hook in hooks:
        model.register_backward_hook(hook)
        
    optim = variable_lr_optimizer()
    stats = defaultdict(list)

    for epoch in range(epochs):
        model.to(device).train()
        e_loss = 0.0
        e_total = 0
#         for i, (x, y) in enumerate(train_dl):
#             optim.zero_grad()
#             x.to(device)
#             y.to(device)
#             bs = x.shape[0]
            
#             # Forward pass
#             y_hat = model(x)
#             loss = criterion(y_hat, y, reduction='mean')
            
#             # Backward pass
#             loss.backward()
#             optim.step()
            
#             # Update mini batch stats.
#             e_total += bs
#             loss += loss * bs
            
#         # Evaluate on validation set.
#         val_stats = validation_metrics()
        
#         # Update epoch stats.
#         stats['loss'].append(e_total)
#         stats['val_loss'].append()
        
        # Print epoch stats.
            
    return stats
def gradient_stats_hook(model, grad_in, grad_out):
    print(stats(grad_out))
train(3, None, None, hnet, nn.BCEWithLogitsLoss, [gradient_stats_hook])<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup