fix f16

app.py

@@ -215,6 +215,11 @@ def run_classifier(image: Image.Image, threshold):
     img = image.convert('RGBA')
     tensor = transform(img).unsqueeze(0)
 
+    if torch.cuda.is_available():
+        tensor = tensor.to(device, dtype=torch.float16)
+    else:
+        tensor = tensor.to(device)
+        
     with torch.no_grad():
         probits = model(tensor)[0] # type: torch.Tensor
         values, indices = probits.cpu().topk(250)
@@ -238,6 +243,13 @@ def cam_inference(img, threshold, alpha, evt: gr.SelectData):
     target_tag_index = tags[evt.value]
     tensor = transform(img).unsqueeze(0)
 
+    if torch.cuda.is_available():
+        tensor = tensor.to(device, dtype=torch.float16)
+    else:
+        tensor = tensor.to(device)
+    
+    tensor.requires_grad_()
+
     gradients = {}
     activations = {}
 
@@ -339,7 +351,10 @@ def process_images(images, threshold):
     all_results = []
     with torch.no_grad():
         for batch, filenames in dataloader:
-            batch = batch.to(device) 
+            if torch.cuda.is_available():
+                batch = batch.to(device, dtype=torch.float16)
+            else:
+                batch = batch.to(device)
             probabilities = model(batch)
             for i, prob in enumerate(probabilities):
                 indices = torch.where(prob > threshold)[0]