diff --git a/checkpoints/ROMA_UNSB_001/loss_log.txt b/checkpoints/ROMA_UNSB_001/loss_log.txt
index a6f9e99..3f966d3 100644
--- a/checkpoints/ROMA_UNSB_001/loss_log.txt
+++ b/checkpoints/ROMA_UNSB_001/loss_log.txt
@@ -1,3 +1,4 @@
 ================ Training Loss (Sun Feb 23 15:46:44 2025) ================
 ================ Training Loss (Sun Feb 23 15:52:29 2025) ================
 ================ Training Loss (Sun Feb 23 16:00:07 2025) ================
+================ Training Loss (Sun Feb 23 16:02:40 2025) ================
diff --git a/checkpoints/ROMA_UNSB_001/train_opt.txt b/checkpoints/ROMA_UNSB_001/train_opt.txt
index 18c3b2c..b15424d 100644
--- a/checkpoints/ROMA_UNSB_001/train_opt.txt
+++ b/checkpoints/ROMA_UNSB_001/train_opt.txt
@@ -63,6 +63,7 @@ nce_includes_all_negatives_from_minibatch: False
                     normG: instance                      
               num_patches: 256                           
               num_threads: 4                             
+            num_timesteps: 10                            
                 output_nc: 3                             
                     phase: train                         
                 pool_size: 0                             
diff --git a/models/__pycache__/roma_unsb_model.cpython-39.pyc b/models/__pycache__/roma_unsb_model.cpython-39.pyc
index 1b4d953..870fbd6 100644
Binary files a/models/__pycache__/roma_unsb_model.cpython-39.pyc and b/models/__pycache__/roma_unsb_model.cpython-39.pyc differ
diff --git a/models/roma_unsb_model.py b/models/roma_unsb_model.py
index 691121d..0c88585 100644
--- a/models/roma_unsb_model.py
+++ b/models/roma_unsb_model.py
@@ -405,6 +405,9 @@ class RomaUnsbModel(BaseModel):
             self.mutil_real_A0_tokens = self.netPreViT(real_A0, self.atten_layers, get_tokens=True)
             self.mutil_real_A1_tokens = self.netPreViT(real_A1, self.atten_layers, get_tokens=True)
 
+        self.mutil_real_A0_tokens = torch.tensor(self.mutil_real_A0_tokens, device=self.device)
+        self.mutil_real_A1_tokens = torch.tensor(self.mutil_real_A1_tokens, device=self.device)
+        
         # 执行一次SB模块
         
         # ============ 第一步：初始化时间步与时间索引 ============