def test_moco(tmpdir):
seed_everything()
datamodule = CIFAR10DataModule(tmpdir, num_workers=0, batch_size=2)
datamodule.train_transforms = Moco2TrainCIFAR10Transforms()
datamodule.val_transforms = Moco2EvalCIFAR10Transforms()
model = MocoV2(data_dir=tmpdir, batch_size=2, online_ft=True)
trainer = pl.Trainer(fast_dev_run=True, max_epochs=1, default_root_dir=tmpdir, callbacks=[MocoLRScheduler()])
trainer.fit(model, datamodule=datamodule)
loss = trainer.progress_bar_dict['loss']
assert float(loss) > 0
def test_moco(tmpdir):
reset_seed()
datamodule = CIFAR10DataModule(tmpdir, num_workers=0)
datamodule.train_transforms = Moco2TrainCIFAR10Transforms()
datamodule.val_transforms = Moco2EvalCIFAR10Transforms()
model = MocoV2(data_dir=tmpdir, batch_size=2, datamodule=datamodule, online_ft=True)
trainer = pl.Trainer(overfit_batches=2, max_epochs=1, default_root_dir=tmpdir, callbacks=[MocoLRScheduler()])
trainer.fit(model)
loss = trainer.callback_metrics['loss']
assert loss > 0