These lines overwrite the device argument of the function, renderin…

…g it useless.
mahmoodlab · Jun 28, 2022 · 73d5c90 · 73d5c90
1 parent aec40d8
commit 73d5c90
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/HIPT_4K/hipt_model_utils.py b/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