From 124f4326d21bd5a8b0df17d9bf73d91461cc5d24 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Thu, 19 Dec 2024 03:08:05 -0500
Subject: [PATCH] Fix setup_seed function to handle non-list seed

---
 deepmd/pt/utils/dataloader.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py
index 907fee7487..6daa1a3435 100644
--- a/deepmd/pt/utils/dataloader.py
+++ b/deepmd/pt/utils/dataloader.py
@@ -52,6 +52,8 @@
 def setup_seed(seed) -> None:
     if isinstance(seed, (list, tuple)):
         mixed_seed = mix_entropy(seed)
+    else:
+        mixed_seed = seed
     torch.manual_seed(mixed_seed)
     torch.cuda.manual_seed_all(mixed_seed)
     torch.backends.cudnn.deterministic = True