ProGamerGov / neural-style-pt

PyTorch implementation of neural style transfer algorithm
MIT License
834 stars 178 forks source link

Support for using the CPU as a device #21

Closed ProGamerGov closed 4 years ago

ProGamerGov commented 4 years ago

For example this put the first layer of the model on the CPU, while every other layer is on the GPU::

python3 neural_style.py -gpu c,0 -image_size 256 -multidevice_strategy 1

And this will put the first and last layers of a model on the CPU, while every other layer is on the GPU:

python3 neural_style.py -gpu c,0,c -image_size 256 -multidevice_strategy 1,34