sr

app.py

@@ -83,11 +83,16 @@ def generate_imgs(x: torch.Tensor,
 
     if hasattr(physics.physics, 'tensor_size'):
         physics.physics.tensor_size = x.shape[1:]
+    elif hasattr(physics.physics, 'imsize'):
+        physics.physics.imsize = 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:]
             physics.generator.tensor_size = x.shape[1:]
+        if hasattr(physics.physics_generator, 'imsize'):
+            physics.physics_generator.imsize = x.shape[1:]
+            physics.generator.imsize = x.shape[1:]
 
     ### Compute y
     with torch.no_grad():