multi-thread context torch.inverse

[wyli]

Traceback (most recent call last):
  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/localhome/Projects/monailabel/monailabel/interfaces/utils/app.py", line 132, in <module>
    run_main()
  File "/localhome/Projects/monailabel/monailabel/interfaces/utils/app.py", line 121, in run_main
    result = a.batch_infer(request)
  File "/localhome/Projects/monailabel/monailabel/interfaces/app.py", line 338, in batch_infer
    return self._batch_infer(request, datastore if datastore else self.datastore(), self.infer)
  File "/localhome/Projects/monailabel/monailabel/interfaces/tasks/batch_infer.py", line 100, in __call__
    res = future.result()
  File "/usr/lib/python3.10/concurrent/futures/_base.py", line 458, in result
    return self.__get_result()
  File "/usr/lib/python3.10/concurrent/futures/_base.py", line 403, in __get_result
    raise self._exception
  File "/usr/lib/python3.10/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/localhome/Projects/monailabel/monailabel/interfaces/tasks/batch_infer.py", line 124, in run_infer_task
    r = infer(req, datastore)
  File "/localhome/Projects/monailabel/monailabel/interfaces/app.py", line 300, in infer
    result_file_name, result_json = task(request)
  File "/localhome/Projects/monailabel/monailabel/tasks/infer/basic_infer.py", line 297, in __call__
    data = self.run_pre_transforms(data, pre_transforms)
  File "/localhome/Projects/monailabel/monailabel/tasks/infer/basic_infer.py", line 388, in run_pre_transforms
    return run_transforms(data, transforms, log_prefix="PRE", use_compose=False)
  File "/localhome/Projects/monailabel/monailabel/interfaces/utils/transform.py", line 106, in run_transforms
    data = t(data)
  File "/localhome/Projects/MONAI/monai/transforms/spatial/dictionary.py", line 416, in __call__
    d[key] = self.spacing_transform(
  File "/localhome/Projects/MONAI/monai/utils/deprecate_utils.py", line 221, in _wrapper
    return func(*args, **kwargs)
  File "/localhome/Projects/MONAI/monai/transforms/spatial/array.py", line 606, in __call__
    data_array = self.sp_resample(
  File "/localhome/Projects/MONAI/monai/utils/deprecate_utils.py", line 221, in _wrapper
    return func(*args, **kwargs)
  File "/localhome/Projects/MONAI/monai/transforms/spatial/array.py", line 315, in __call__
    img = affine_xform(img.unsqueeze(0), theta=xform, spatial_size=spatial_size).squeeze(0)
  File "/localhome/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/localhome/Projects/MONAI/monai/networks/layers/spatial_transforms.py", line 556, in forward
    theta = to_norm_affine(
  File "/localhome/Projects/MONAI/monai/networks/utils.py", line 279, in to_norm_affine
    return src_xform @ affine @ torch.inverse(dst_xform)
RuntimeError: lazy wrapper should be called at most once

[wyli]

bug report pytorch/pytorch#90613

[wyli]

this is reproducible with

from __future__ import annotations

import concurrent.futures

import torch

import monai.transforms as mt
from monai.data import MetaTensor


def run_task(req):
    x = MetaTensor(torch.zeros(1, 3, 3, 3).to("cuda:0"))
    out = mt.Spacing(pixdim=1.5)(x)
    return req, out


infer_tasks = [1, 2, 3, 4, 5, 6, 7, 8]

with concurrent.futures.ThreadPoolExecutor(4, "BatchTask") as executor:
    futures = {executor.submit(run_task, t) for t in infer_tasks}
    for fut in concurrent.futures.as_completed(futures):
        print("outcome", fut.result())

Traceback (most recent call last):
  File "my_test_6.py", line 22, in <module>
    print("outcome", fut.result())
  File "/home/anaconda3/envs/py38/lib/python3.8/concurrent/futures/_base.py", line 437, in result
    return self.__get_result()
  File "/home/anaconda3/envs/py38/lib/python3.8/concurrent/futures/_base.py", line 389, in __get_result
    raise self._exception
  File "/home/anaconda3/envs/py38/lib/python3.8/concurrent/futures/thread.py", line 57, in run
    result = self.fn(*self.args, **self.kwargs)
  File "my_test_6.py", line 13, in run_task
    out = mt.Spacing(pixdim=1.5)(x)
  File "/home/Documents/MONAI/monai/utils/deprecate_utils.py", line 223, in _wrapper
    return func(*args, **kwargs)
  File "/home/Documents/MONAI/monai/transforms/spatial/array.py", line 586, in __call__
    data_array = self.sp_resample(
  File "/home/Documents/MONAI/monai/transforms/spatial/array.py", line 315, in __call__
    img = affine_xform(img.unsqueeze(0), theta=xform, spatial_size=spatial_size).squeeze(0)  # type: ignore
  File "/home/anaconda3/envs/py38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/Documents/MONAI/monai/networks/layers/spatial_transforms.py", line 558, in forward
    theta = to_norm_affine(
  File "/home/Documents/MONAI/monai/networks/utils.py", line 268, in to_norm_affine
    return src_xform @ affine @ torch.inverse(dst_xform)
RuntimeError: lazy wrapper should be called at most once

Numpy version: 1.23.5
Pytorch version: 1.13.1+cu117