Method get_model_optimizer

tests/others/test_training.py:30–34

(self, resolution=32)

class TrainingTests(unittest.TestCase):
    def get_model_optimizer(self, resolution=32):
        set_seed(0)
        model = UNet2DModel(sample_size=resolution, in_channels=3, out_channels=3)
        optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)
        return model, optimizer

    @slow
    def test_training_step_equality(self):