These lines overwrite the device argument of the function, rendering it useless.

Author: jjhbw

HIPT_4K/hipt_4k.py

@@ -31,8 +31,8 @@ def __init__(self,
 		device4k=torch.device('cuda:1')):
 
 		super().__init__()
-		self.model256 = get_vit256(pretrained_weights=model256_path).to(device256)
-		self.model4k = get_vit4k(pretrained_weights=model4k_path).to(device4k)
+		self.model256 = get_vit256(pretrained_weights=model256_path, device=device256).to(device256)
+		self.model4k = get_vit4k(pretrained_weights=model4k_path, device=device4k).to(device4k)
 		self.device256 = device256
 		self.device4k = device4k
 

HIPT_4K/hipt_model_utils.py

@@ -32,7 +32,6 @@ def get_vit256(pretrained_weights, arch='vit_small', device=torch.device('cuda:0
     """
 
     checkpoint_key = 'teacher'
-    device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
     model256 = vits.__dict__[arch](patch_size=16, num_classes=0)
     for p in model256.parameters():
         p.requires_grad = False
@@ -68,7 +67,6 @@ def get_vit4k(pretrained_weights, arch='vit4k_xs', device=torch.device('cuda:1')
     """
 
     checkpoint_key = 'teacher'
-    device = torch.device("cuda:1") if torch.cuda.is_available() else torch.device("cpu")
     model4k = vits4k.__dict__[arch](num_classes=0)
     for p in model4k.parameters():
         p.requires_grad = False