diff --git a/checkpoints/ROMA_UNSB_001/loss_log.txt b/checkpoints/ROMA_UNSB_001/loss_log.txt
index 95fd4b1..09b323f 100644
--- a/checkpoints/ROMA_UNSB_001/loss_log.txt
+++ b/checkpoints/ROMA_UNSB_001/loss_log.txt
@@ -4,3 +4,4 @@
 ================ Training Loss (Sun Feb 23 16:02:40 2025) ================
 ================ Training Loss (Sun Feb 23 16:05:19 2025) ================
 ================ Training Loss (Sun Feb 23 16:06:44 2025) ================
+================ Training Loss (Sun Feb 23 16:09:38 2025) ================
diff --git a/models/__pycache__/roma_unsb_model.cpython-39.pyc b/models/__pycache__/roma_unsb_model.cpython-39.pyc
index eb910dd..f14a8b5 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/networks.py b/models/networks.py
index 933f792..ebae8aa 100644
--- a/models/networks.py
+++ b/models/networks.py
@@ -984,7 +984,7 @@ class ResnetGenerator(nn.Module):
                 else:
                     # print("%d: skipping %s %d" % (layer_id, layer.__class__.__name__, feat.size(1)))
                     pass
-                if layer_id == layers[-1] and encode_only:
+                if layer_id == len(list(self.model)) - 1 and encode_only:
                     # print('encoder only return features')
                     return feats  # return intermediate features alone; stop in the last layers
 
diff --git a/models/roma_unsb_model.py b/models/roma_unsb_model.py
index 0000f0f..11ce30c 100644
--- a/models/roma_unsb_model.py
+++ b/models/roma_unsb_model.py
@@ -405,7 +405,6 @@ 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)
 
-        print(self.mutil_real_A0_tokens)
         self.mutil_real_A0_tokens = torch.cat(self.mutil_real_A0_tokens, dim=1).to(self.device)
         self.mutil_real_A1_tokens = torch.cat(self.mutil_real_A1_tokens, dim=1).to(self.device)