NoCrypt commited on
Commit
791b9ee
1 Parent(s): b45103b

Update models/c2pGen.py

Browse files
Files changed (1) hide show
  1. 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