Update app.py

app.py

@@ -59,9 +59,9 @@ def reply(modelA, modelB, prompt):
     df.loc[df["MODEL"] == modelA, "MATCHES_PLAYED"] += 1
     df.loc[df["MODEL"] == modelB, "MATCHES_PLAYED"] += 1
     df.to_csv("models.csv", index=False)
-    pipeA = pipeline("text-generation", model=models_and_tokenizers[modelA][0], tokenizer=models_and_tokenizers[modelA][1], max_new_tokens=512, repetition_penalty=1.5, temperature=0.5, device_map="auto")
+    pipeA = pipeline("text-generation", model=models_and_tokenizers[modelA][0], tokenizer=models_and_tokenizers[modelA][1], max_new_tokens=512, repetition_penalty=1.5, temperature=0.5, device_map="cuda:0")
     responseA = run_inference(pipeA, prompt)    
-    pipeB = pipeline("text-generation", model=models_and_tokenizers[modelB][0], tokenizer=models_and_tokenizers[modelB][1], max_new_tokens=512, repetition_penalty=1.5, temperature=0.5, device_map="auto")
+    pipeB = pipeline("text-generation", model=models_and_tokenizers[modelB][0], tokenizer=models_and_tokenizers[modelB][1], max_new_tokens=512, repetition_penalty=1.5, temperature=0.5, device_map="cuda:1")
     responseB = run_inference(pipeB, prompt)    
     return responseA, responseB