Update fs_model.py

This commit is contained in:
Smiril
2022-07-12 12:31:05 +02:00
committed by GitHub
parent 8e616d1141
commit c01d582b3f

View File

@@ -48,7 +48,13 @@ class fsModel(BaseModel):
torch.backends.cudnn.benchmark = True
self.isTrain = opt.isTrain
device = torch.device("cuda:0")
if not torch.backends.mps.is_available():
if not torch.backends.mps.is_built():
device = torch.device("cuda:0")
else:
print("ERROR")
else:
device = torch.device("mps")
if opt.crop_size == 224:
from .fs_networks import Generator_Adain_Upsample, Discriminator