def test_forward(net, img_size):
    """Test `.forward()` doesn't throw an error"""
    data = torch.zeros((1, 3, img_size, img_size))
    output = net(data)
    assert not torch.isnan(output).any()
Test `.forward()` doesn't throw an error
tests/test_model.py