Update app.py

app.py

@@ -439,7 +439,7 @@ class UniMCPredict:
         batch = [self.data_model.train_data.encode(
             sample) for sample in batch_data]
         batch = self.data_model.collate_fn(batch)
-        batch = {k: v.cuda() for k, v in batch.items()}
+        # batch = {k: v.cuda() for k, v in batch.items()}
         _, _, logits = self.model.model(**batch)
         soft_logits = torch.nn.functional.softmax(logits, dim=-1)
         logits = torch.argmax(soft_logits, dim=-1).detach().cpu().numpy()