From 065ddaa77be5184530ffcdff1baf99efab9fda27 Mon Sep 17 00:00:00 2001
From: Hamid Shojanazeri <hamid.nazeri2010@gmail.com>
Date: Sat, 22 Jul 2023 21:18:29 +0000
Subject: [PATCH] fixing the condition for moving to cuda
---
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 c6745d5c..defb6518 100644
--- a/utils/train_utils.py
+++ b/utils/train_utils.py
@@ -85,7 +85,7 @@ def train(model, train_dataloader,eval_dataloader, tokenizer, optimizer, lr_sche
for key in batch.keys():
if train_config.enable_fsdp:
batch[key] = batch[key].to(local_rank)
- elif not train_config.quantization:
+ else:
batch[key] = batch[key].to('cuda')
outputs = model(**batch)
loss = outputs.loss
--
GitLab