Update README

README.md

@@ -1,8 +1,8 @@
 ---
 title: TCL
 emoji: 🍁
-colorFrom: teal
-colorTo: teal
+colorFrom: yellow
+colorTo: red
 sdk: gradio
 sdk_version: 3.23.0
 app_file: app.py

app.py

@@ -23,7 +23,6 @@ import json
 from contextlib import ExitStack
 import gradio as gr
 import torch
-from torch.cuda.amp import autocast
 
 from detectron2.evaluation import inference_context
 
@@ -117,8 +116,10 @@ def inference(img, query):
         stack.enter_context(inference_context(model))
         stack.enter_context(torch.no_grad())
 
-        with autocast():
-            visualized_output = model.forward_vis(img, query)
+        if device.type == "cuda":
+            stack.enter_context(torch.autocast("cuda"))
+
+        visualized_output = model.forward_vis(img, query)
 
         return visualized_output