inpainting debug

app.py

@@ -80,6 +80,17 @@ def generate_imgs(x: torch.Tensor,
     print(f"[Before inference] CUDA max allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
     print(f"[Before inference] CUDA max reserved: {torch.cuda.max_memory_reserved() / 1024**2:.2f} MB")
 
+    if hasattr(physics.physics, 'imsize'):
+        physics.physics.imsize = x.shape[1:]
+    elif hasattr(physics.physics, 'img_size'):
+        physics.physics.img_size = x.shape[1:]
+    elif hasattr(physics.physics, 'tensor_size'):
+        physics.physics.tensor_size = x.shape[1:]
+
+    if physics.physics_generator is not None:  # we only change physic params but not noise levels
+        if hasattr(physics.physics_generator, 'tensor_size'):
+            physics.physics_generator.tensor_size = x.shape[1:]
+
     ### Compute y
     with torch.no_grad():
         y = physics(x, use_gen)  # possible reduction in img shape due to Blurring