Move queue to gpu when resuming checkpoint - SWAV self supervised model

CHANGELOG.md

@@ -32,6 +32,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Removed momentum updating from val step and add separate val queue ([#631](https://github.com/PyTorchLightning/lightning-bolts/pull/631))
 
+- Fixed moving the queue to GPU when resuming checkpoint for SwAV model ([#684](https://github.com/PyTorchLightning/lightning-bolts/pull/684))
+
 
 - Fixed FP16 support with vision GPT model ([#694](https://github.com/PyTorchLightning/lightning-bolts/pull/694))
 

pl_bolts/models/self_supervised/swav/swav_module.py

@@ -183,8 +183,8 @@ def on_train_epoch_start(self):
  self.feat_dim,
  )
 
-  if self.gpus > 0:
-  self.queue = self.queue.cuda()
+ if self.queue is not None:
+ self.queue = self.queue.to(self.device)
 
  self.use_the_queue = False