Move model's object to a correct device (#172)

Co-authored-by: Ilya Isaev <[email protected]>

Author: IsaevIlya

s3torchbenchmarking/src/s3torchbenchmarking/models.py

@@ -102,7 +102,7 @@ def __init__(self, num_labels: int, checkpoint: DictConfig):
     def model(self):
         return ViTForImageClassification.from_pretrained(
             "google/vit-base-patch16-224-in21k", num_labels=self.num_labels
-        )
+        ).to(self.device)
 
     @cached_property
     def transform(self):