Lungmask version: 0.2.18 throws exception in 3D Slicer

[rbumm]

Lungmask version: 0.2.18
[Python] Failed to compute results: CUDA error: the launch timed out and was terminated
[Python] CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
[Python] For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
[Python] Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
Traceback (most recent call last):
  File "C:/Users/rudol/AppData/Local/slicer.org/Slicer 5.6.0/slicer.org/Extensions-32390/LungCTAnalyzer/lib/Slicer-5.6/qt-scripted-modules/LungCTSegmenter.py", line 1014, in runProcessing
    self.logic.applySegmentation()
  File "C:/Users/rudol/AppData/Local/slicer.org/Slicer 5.6.0/slicer.org/Extensions-32390/LungCTAnalyzer/lib/Slicer-5.6/qt-scripted-modules/LungCTSegmenter.py", line 2362, in applySegmentation
    inferer = LMInferer()
  File "C:\Users\rudol\AppData\Local\slicer.org\Slicer 5.6.0\lib\Python\Lib\site-packages\lungmask\mask.py", line 137, in __init__
    self.model.to(self.device)
  File "C:\Users\rudol\AppData\Local\slicer.org\Slicer 5.6.0\lib\Python\Lib\site-packages\torch\nn\modules\module.py", line 1160, in to
    return self._apply(convert)
  File "C:\Users\rudol\AppData\Local\slicer.org\Slicer 5.6.0\lib\Python\Lib\site-packages\torch\nn\modules\module.py", line 810, in _apply
    module._apply(fn)
  File "C:\Users\rudol\AppData\Local\slicer.org\Slicer 5.6.0\lib\Python\Lib\site-packages\torch\nn\modules\module.py", line 810, in _apply
    module._apply(fn)
  File "C:\Users\rudol\AppData\Local\slicer.org\Slicer 5.6.0\lib\Python\Lib\site-packages\torch\nn\modules\module.py", line 810, in _apply
    module._apply(fn)
  File "C:\Users\rudol\AppData\Local\slicer.org\Slicer 5.6.0\lib\Python\Lib\site-packages\torch\nn\modules\module.py", line 833, in _apply
    param_applied = fn(param)
  File "C:\Users\rudol\AppData\Local\slicer.org\Slicer 5.6.0\lib\Python\Lib\site-packages\torch\nn\modules\module.py", line 1158, in convert
    return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
RuntimeError: CUDA error: the launch timed out and was terminated
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

[VTK] Warning: In vtkMRMLSegmentationDisplayNode.cxx, line 295
[VTK] vtkMRMLSegmentationDisplayNode (0000029F957C0D40): vtkMRMLSegmentationDisplayNode::GetSegmentDisplayProperties: no display properties are found for segment ID=, return default

Installation is done:

lungmaskPackage = "https://github.com/JoHof/lungmask/archive/master.zip"
slicer.util.pip_install(lungmaskPackage)
import lungmask

[rbumm]

Driver error. Updating the NVIDIA Driver and Pytorch resolved the problem.