github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 06:58:11 -05:00
Code Issues Packages Projects Releases Wiki Activity

hotfix: resnet to obj.device

Browse Source
This commit is contained in:
George Hotz
2024-09-06 13:06:02 +08:00
parent 9d72119a0c
commit 8f6d0485e7
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
extra/models/resnet.py
View File

@@ -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)