From 8f6d0485e7fe3f68bfa32c28b6caf0288be652a9 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 6 Sep 2024 13:06:02 +0800 Subject: [PATCH] hotfix: resnet to obj.device --- extra/models/resnet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/extra/models/resnet.py b/extra/models/resnet.py index dc07861ea9..9a4120476a 100644 --- a/extra/models/resnet.py +++ b/extra/models/resnet.py @@ -144,7 +144,7 @@ class ResNet: continue # Skip FC if transfer learning if 'bn' not in k and 'downsample' not in k: assert obj.shape == dat.shape, (k, obj.shape, dat.shape) - obj.assign(dat.to(None).reshape(obj.shape)) + obj.assign(dat.to(obj.device).reshape(obj.shape)) ResNet18 = lambda num_classes=1000: ResNet(18, num_classes=num_classes) ResNet34 = lambda num_classes=1000: ResNet(34, num_classes=num_classes)