diff --git a/inference_mp4_2x.py b/inference_mp4_2x.py
index 01a0b68..bfe7fda 100644
--- a/inference_mp4_2x.py
+++ b/inference_mp4_2x.py
@@ -48,9 +48,9 @@ if args.montage:
 while success:
     lastframe = frame
     success, frame = videoCapture.read()
-    if args.montage:
-        frame = frame[:, left: left + w]
     if success:
+        if args.montage:
+            frame = frame[:, left: left + w]if args.montage:            
         I0 = torch.from_numpy(np.transpose(lastframe, (2,0,1)).astype("float32") / 255.).to(device).unsqueeze(0)
         I1 = torch.from_numpy(np.transpose(frame, (2,0,1)).astype("float32") / 255.).to(device).unsqueeze(0)
         I0 = F.pad(I0, padding)
diff --git a/inference_mp4_4x.py b/inference_mp4_4x.py
index f1b2624..20dd593 100644
--- a/inference_mp4_4x.py
+++ b/inference_mp4_4x.py
@@ -48,9 +48,9 @@ if args.montage:
 while success:
     lastframe = frame
     success, frame = videoCapture.read()
-    if args.montage:
-        frame = frame[:, left: left + w]
     if success:
+        if args.montage:
+            frame = frame[:, left: left + w]
         I0 = torch.from_numpy(np.transpose(lastframe, (2,0,1)).astype("float32") / 255.).to(device).unsqueeze(0)
         I1 = torch.from_numpy(np.transpose(frame, (2,0,1)).astype("float32") / 255.).to(device).unsqueeze(0)
         I0 = F.pad(I0, padding)