small compatibility fixes (#220)

vturrisi · Jan 25, 2022 · e8e2303 · e8e2303
1 parent d5b96da
commit e8e2303
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/solo/args/utils.py b/solo/args/utils.py
@@ -251,12 +251,10 @@ def additional_setup_linear(args: Namespace):
 
     # create backbone-specific arguments
     args.backbone_args = {"cifar": args.dataset in ["cifar10", "cifar100"]}
-
     if "resnet" not in args.backbone:
         # dataset related for all transformers
         crop_size = args.crop_size[0]
         args.backbone_args["img_size"] = crop_size
-
         if "vit" in args.backbone:
             args.backbone_args["patch_size"] = args.patch_size
 

diff --git a/tests/utils/test_checkpointer.py b/tests/utils/test_checkpointer.py
@@ -91,7 +91,8 @@ def test_checkpointer():
         "optimizer_states",
         "lr_schedulers",
     ]
-    assert list(ckpt.keys()) == expected_keys
+    ckpt_keys = list(ckpt.keys())
+    assert all(k in ckpt_keys for k in expected_keys)
 
     parser = argparse.ArgumentParser()
     ckpt_callback.add_checkpointer_args(parser)