From a7156dfb5d0bcc01f52c82ac03b5187797a07133 Mon Sep 17 00:00:00 2001
From: Hamid Shojanazeri <hamid.nazeri2010@gmail.com>
Date: Tue, 25 Jul 2023 06:48:17 +0000
Subject: [PATCH] fixing the cuda id
---
utils/train_utils.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/utils/train_utils.py b/utils/train_utils.py
index 8a68f0c1..e41f503e 100644
--- a/utils/train_utils.py
+++ b/utils/train_utils.py
@@ -199,7 +199,7 @@ def evaluation(model,train_config, eval_dataloader, local_rank, tokenizer):
if train_config.enable_fsdp:
batch[key] = batch[key].to(local_rank)
else:
- batch[key] = batch[key].to('cuda')
+ batch[key] = batch[key].to('cuda:0')
# Ensure no gradients are computed for this scope to save memory
with torch.no_grad():
# Forward pass and compute loss
--
GitLab