Update model/model_manager.py

model/model_manager.py

@@ -39,7 +39,7 @@ class ModelManager:
     @spaces.GPU(duration=30)
     def NSFW_filter(self, prompt):
         chat = [{"role": "user", "content": prompt}]
-        input_ids = self.tokenizer.apply_chat_template(chat, return_tensors="pt").to(device)
+        input_ids = self.tokenizer.apply_chat_template(chat, return_tensors="pt").to('cuda')
         self.guard.cuda()
         output = self.guard.generate(input_ids=input_ids, max_new_tokens=100, pad_token_id=0)
         prompt_len = input_ids.shape[-1]