stylegan2-pytorch RuntimeError: Unknown type name 'torch.device':

`File "/home/lbc/.local/bin/stylegan2_pytorch", line 5, in from stylegan2_pytorch.cli import main File "/home/lbc/.local/lib/python3.6/site-packages/stylegan2_pytorch/init.py", line 1, in from stylegan2_pytorch.stylegan2_pytorch import Trainer, StyleGAN2, NanException File "/home/lbc/.local/lib/python3.6/site-packages/stylegan2_pytorch/stylegan2_pytorch.py", line 26, in from kornia.filters import filter2D File "/home/lbc/.local/lib/python3.6/site-packages/kornia/init.py", line 19, in from kornia import jit File "/home/lbc/.local/lib/python3.6/site-packages/kornia/jit/init.py", line 8, in spatial_soft_argmax2d = torch.jit.script(K.geometry.spatial_soft_argmax2d) File "/home/lbc/.local/lib/python3.6/site-packages/torch/jit/init.py", line 1281, in script fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj)) File "/home/lbc/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 555, in try_compile_fn return torch.jit.script(fn, _rcb=rcb) File "/home/lbc/.local/lib/python3.6/site-packages/torch/jit/init.py", line 1281, in script fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj)) File "/home/lbc/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 555, in try_compile_fn return torch.jit.script(fn, _rcb=rcb) File "/home/lbc/.local/lib/python3.6/site-packages/torch/jit/init.py", line 1281, in script fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj)) RuntimeError: Unknown type name 'torch.device': File "/home/lbc/.local/lib/python3.6/site-packages/kornia/utils/grid.py", line 10 width: int, normalized_coordinates: bool = True, device: Optional[torch.device] = torch.device('cpu')) -> torch.Tensor: ~~~~~~~~~~~~ <--- HERE """Generates a coordinate grid for an image. 'create_meshgrid' is being compiled since it was called from 'spatial_expectation2d' File "/home/lbc/.local/lib/python3.6/site-packages/kornia/geometry/subpix/dsnt.py", line 89

# Create coordinates grid.
grid: torch.Tensor = create_meshgrid(height, width, normalized_coordinates, input.device)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
grid = grid.to(input.dtype)

'spatial_expectation2d' is being compiled since it was called from 'spatial_soft_argmax2d' File "/home/lbc/.local/lib/python3.6/site-packages/kornia/geometry/subpix/spatial_soft_argmax.py", line 516 """ input_soft: torch.Tensor = dsnt.spatial_softmax2d(input, temperature) output: torch.Tensor = dsnt.spatial_expectation2d(input_soft, normalized_coordinates) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE return output`