diff --git a/checkpoints/ROMA_UNSB_001/loss_log.txt b/checkpoints/ROMA_UNSB_001/loss_log.txt
index 485f77f..a6f9e99 100644
--- a/checkpoints/ROMA_UNSB_001/loss_log.txt
+++ b/checkpoints/ROMA_UNSB_001/loss_log.txt
@@ -1,2 +1,3 @@
 ================ 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) ================
diff --git a/checkpoints/ROMA_UNSB_001/train_opt.txt b/checkpoints/ROMA_UNSB_001/train_opt.txt
index 0658b40..18c3b2c 100644
--- a/checkpoints/ROMA_UNSB_001/train_opt.txt
+++ b/checkpoints/ROMA_UNSB_001/train_opt.txt
@@ -76,6 +76,7 @@ nce_includes_all_negatives_from_minibatch: False
            serial_batches: False                         
 stylegan2_G_num_downsampling: 1                             
                    suffix:                               
+                      tau: 0.1                           
          update_html_freq: 1000                          
                   use_idt: False                         
                   verbose: False                         
diff --git a/models/__pycache__/roma_unsb_model.cpython-39.pyc b/models/__pycache__/roma_unsb_model.cpython-39.pyc
index 277be8f..1b4d953 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 4efe49e..691121d 100644
--- a/models/roma_unsb_model.py
+++ b/models/roma_unsb_model.py
@@ -222,6 +222,7 @@ class RomaUnsbModel(BaseModel):
         parser.add_argument('--atten_layers', type=str, default='1,3,5', help='compute Cross-Similarity on which layers')
         
         parser.add_argument('--tau', type=float, default=0.1, help='used in unsb')
+        parser.add_argument('--num_timesteps', type=int, default=10, help='used in unsb')
 
         parser.set_defaults(pool_size=0)  # no image pooling
 
diff --git a/scripts/train.sh b/scripts/train.sh
index c0a4781..d005fe5 100755
--- a/scripts/train.sh
+++ b/scripts/train.sh
@@ -28,4 +28,5 @@ python train.py \
     --num_patches 256 \
     --flip_equivariance False \
     --eta_ratio 0.1 \
-    --tau 0.1
+    --tau 0.1 \
+    --num_timesteps 10