Spaces:

pedrogengo
/

style_loss_showdown

Runtime error

pedrogengo commited on Mar 3, 2023

Commit

ebe933f

1 Parent(s): 8e1cf8c

Update main.py

Files changed (1) hide show

main.py CHANGED Viewed

@@ -216,7 +216,7 @@ map_style_losses = {
 if __name__ == "__main__":
-  device = 'cuda' if torch.cuda.is_available else 'cpu'
   vgg16 = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).to(device)
   vgg16.eval()
   vgg16 = vgg16.features

 if __name__ == "__main__":
+  device = 'cuda' if torch.cuda.is_available() else 'cpu'
   vgg16 = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).to(device)
   vgg16.eval()
   vgg16 = vgg16.features