RedRocket
/

JointTaggerProject

Image Classification

Model card Files Files and versions Community

drhead commited on Jun 26

Commit

2c1eb35

•

1 Parent(s): 5cb75da

Update inference_gradio.py

Files changed (1) hide show

inference_gradio.py +1 -1

inference_gradio.py CHANGED Viewed

@@ -137,7 +137,7 @@ def create_tags(image, threshold):
     tensor = transform(img).unsqueeze(0) # type: torch.Tensor
     if torch.cuda.is_available():
-        tensor.cuda()
         if torch.cuda.get_device_capability()[0] >= 7:
             tensor = tensor.to(dtype=torch.float16, memory_format=torch.channels_last)

     tensor = transform(img).unsqueeze(0) # type: torch.Tensor
     if torch.cuda.is_available():
+        tensor = tensor.cuda()
         if torch.cuda.get_device_capability()[0] >= 7:
             tensor = tensor.to(dtype=torch.float16, memory_format=torch.channels_last)