From 639fa264f74cdb061dd4beee480b7f9126ab3f05 Mon Sep 17 00:00:00 2001 From: ProGamerGov Date: Sat, 10 Sep 2022 09:59:43 -0600 Subject: [PATCH] Update neural_style.py --- neural_style.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/neural_style.py b/neural_style.py index 1e0162d..a3a7c2c 100644 --- a/neural_style.py +++ b/neural_style.py @@ -55,7 +55,7 @@ def main(): - dtype, multidevice, backward_device = setup_gpu() + multidevice, backward_device = setup_gpu() cnn, layerList = loadCaffemodel(params.model_file, params.pooling, params.disable_check) @@ -318,18 +318,17 @@ def setup_mps(): else: backward_device = "cuda:" + devices[0] setup_cuda() - dtype = torch.FloatTensor elif "cpu" not in str(params.gpu).lower() and "mps" not in str(params.gpu).lower(): setup_cuda() - dtype, backward_device = torch.cuda.FloatTensor, "cuda:" + str(params.gpu) + backward_device = "cuda:" + str(params.gpu) elif "mps" in str(params.gpu).lower(): setup_mps() - dtype, backward_device = torch.FloatTensor, "mps" + backward_device = "mps" else: setup_cpu() - dtype, backward_device = torch.FloatTensor, "cpu" - return dtype, multidevice, backward_device + backward_device = "cpu" + return multidevice, backward_device def setup_multi_device(net, device, multidevice_strategy):