Upload model

modeling.py

@@ -93,11 +93,12 @@ class BoneAgeCropModel(PreTrainedModel):
         # if img_shape is provided, will provide rescaled coordinates
         # otherwise, provide normalized [0, 1] coordinates
         # coords format is xywh
-        assert x.size(0) == img_shape.size(
-            0
-        ), f"x.size(0) [{x.size(0)}] must equal img_shape.size(0) [{img_shape.size(0)}]"
-        # img_shape = (batch_dim, 2)
-        # img_shape[:, 0] = height, img_shape[:, 1] = width
+        if img_shape is not None:
+            assert x.size(0) == img_shape.size(
+                0
+            ), f"x.size(0) [{x.size(0)}] must equal img_shape.size(0) [{img_shape.size(0)}]"
+            # img_shape = (batch_dim, 2)
+            # img_shape[:, 0] = height, img_shape[:, 1] = width
 
         x = self.normalize(x)
         features = self.pooling(self.backbone(x))