update resnet.load_from_pretrained (#5040)

extra/models/resnet.py

@@ -140,14 +140,16 @@ class ResNet:
     self.url = model_urls[(self.num, self.groups, self.base_width)]
     for k, v in torch_load(fetch(self.url)).items():
       obj: Tensor = get_child(self, k)
-      dat = v.detach().numpy()
+      dat = v.numpy()
 
       if 'fc.' in k and obj.shape != dat.shape:
         print("skipping fully connected layer")
         continue # Skip FC if transfer learning
 
-      # TODO: remove or when #777 is merged
-      assert obj.shape == dat.shape or (obj.shape == (1,) and dat.shape == ()), (k, obj.shape, dat.shape)
+      if dat.shape == ():
+        assert obj.shape == (1,), obj.shape
+        dat = dat.reshape(1)
+      assert obj.shape == dat.shape, (k, obj.shape, dat.shape)
       obj.assign(dat)
 
 ResNet18 = lambda num_classes=1000: ResNet(18, num_classes=num_classes)

test/models/test_resnet.py

@@ -0,0 +1,14 @@
+import unittest
+from extra.models import resnet
+
+class TestResnet(unittest.TestCase):
+  def test_model_load(self):
+    model = resnet.ResNet18()
+    model.load_from_pretrained()
+
+    model = resnet.ResNeXt50_32X4D()
+    model.load_from_pretrained()
+
+
+if __name__ == '__main__':
+  unittest.main()