Module neuralops.tests.test_mixed_precision_epoch

Tests for complete mixed precision trainer with gradient accumulation.

Functions

def test_mixed_precision_train_epoch_grad_accum()

Expand source code

def test_mixed_precision_train_epoch_grad_accum():
    model = nn.Linear(16, 4)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()
    trainer = MixedPrecisionTrainer(model, optimizer, loss_fn, enabled=False)

    data = [(torch.randn(4, 16), torch.randn(4, 4)) for _ in range(6)]
    avg_loss = trainer.train_epoch(data, grad_accum=2, max_steps=4)

    assert isinstance(avg_loss, float)
    assert avg_loss >= 0.0