update SwAV example (#960)

Co-authored-by: Jirka Borovec <[email protected]> Co-authored-by: Jirka B <[email protected]>
Lightning-Universe · May 19, 2023 · ab9e2e7 · ab9e2e7
1 parent c5137b4
commit ab9e2e7
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/docs/source/models/self_supervised.rst b/docs/source/models/self_supervised.rst
@@ -489,12 +489,13 @@ To Train::
         gpus=1,
         num_samples=dm.num_unlabeled_samples,
         dataset='stl10',
-        batch_size=batch_size
+        batch_size=batch_size,
+        nmb_crops=(2,4)
     )
 
     # fit
-    trainer = pl.Trainer(precision=16)
-    trainer.fit(model)
+    trainer = pl.Trainer(precision=16, accelerator='auto')
+    trainer.fit(model, datamodule=dm)
 
 Pre-trained ImageNet
 ********************