Loading styleGAN2 for inference without GPU

I have trained a segmentation model (StyleGAN2) on my machine and want to use the inference in another machine that has not any GPU device.
I tried to modify all imports to use CPU, but still I cannot load it on CPU.

Traceback (most recent call last):
  File "/home/Desktop/SemGAN/inference.py", line 5, in <module>
    from models.encoder_model import ResEncoder
  File "/home/Desktop/SemGAN/models/encoder_model.py", line 27, in <module>
    from models.stylegan2_seg import EqualLinear, ConvLayer, ResBlock
  File "/home/Desktop/SemGAN/models/stylegan2_seg.py", line 11, in <module>
    from models.utils import *
  File "/home/Desktop/SemGAN/models/utils.py", line 12, in <module>
    from models.op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d
  File "/home/Desktop/SemGAN/models/op/__init__.py", line 1, in <module>
    from .fused_act import FusedLeakyReLU, fused_leaky_relu
  File "/home/Desktop/SemGAN/models/op/fused_act.py", line 17, in <module>
    fused = load(
  File "/home/Desktop/land_monitoring/venv/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1202, in load
    return _jit_compile(
  File "/home/Desktop/land_monitoring/venv/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1425, in _jit_compile
    _write_ninja_file_and_build_library(
  File "/home/Desktop/land_monitoring/venv/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1514, in _write_ninja_file_and_build_library
    extra_ldflags = _prepare_ldflags(
  File "/home/Desktop/land_monitoring/venv/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1622, in _prepare_ldflags
    extra_ldflags.append(f'-L{_join_cuda_home("lib64")}')
  File "/home/Desktop/land_monitoring/venv/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 2125, in _join_cuda_home
    raise EnvironmentError('CUDA_HOME environment variable is not set. '
OSError: CUDA_HOME environment variable is not set. Please set it to your CUDA install root.

1 Like

Based on the error message it seems a custom CUDAExtension is supposed to be built, which is failing in your setup. You could check which layer is using this extension and then check if you could replace it with a CPU-only version.