Spaces:
Running
Running
Update models/c2pGen.py
Browse files- models/c2pGen.py +1 -1
models/c2pGen.py
CHANGED
@@ -122,7 +122,7 @@ class PixelBlockEncoder(nn.Module):
|
|
122 |
super(PixelBlockEncoder, self).__init__()
|
123 |
vgg19 = models.vgg.vgg19()
|
124 |
vgg19.classifier._modules['6'] = nn.Linear(4096, 7, bias=True)
|
125 |
-
vgg19.load_state_dict(torch.load('./pixelart_vgg19.pth' if not os.environ['PIX_MODEL'] else os.environ['PIX_MODEL'], map_location=torch.device('cpu')))
|
126 |
self.vgg = vgg19.features
|
127 |
for p in self.vgg.parameters():
|
128 |
p.requires_grad = False
|
|
|
122 |
super(PixelBlockEncoder, self).__init__()
|
123 |
vgg19 = models.vgg.vgg19()
|
124 |
vgg19.classifier._modules['6'] = nn.Linear(4096, 7, bias=True)
|
125 |
+
vgg19.load_state_dict(torch.load('./pixelart_vgg19.pth' if not os.environ['PIX_MODEL'] else os.environ['PIX_MODEL'], map_location=torch.device('cpu'), weights_only=True))
|
126 |
self.vgg = vgg19.features
|
127 |
for p in self.vgg.parameters():
|
128 |
p.requires_grad = False
|