Untitled

import torch
from torch.autograd import gradcheck

# Define your model and loss function
model = YourModel()
loss_fn = YourLossFunction()

# Define your optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# Training loop
for epoch in range(num_epochs):
    for batch_idx, (inputs, targets) in enumerate(train_loader):
        # Forward pass
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

        # Backward pass
        optimizer.zero_grad()
        loss.backward()

        # Check gradients
        inputs.requires_grad_()
        assert gradcheck(loss_fn, inputs), "Gradient check failed!"

        # Update parameters
        optimizer.step()

        # Print loss and other metrics
        if batch_idx % print_interval == 0:
            print("Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}"
                  .format(epoch+1, num_epochs, batch_idx+1, len(train_loader), loss.item()))