PyTorch model failing on CPU

[michaelHampel]

Hello,

great project, thanks for making it available.

For me, running play_amo.py on macOS, MacBookPro m3 fails with:

File "play_amo.py", line 305, in run
    raw_action = self.policy_jit(obs_tensor, extra_hist).cpu().numpy().squeeze()
  File "/Users/michham/github/AMO/.venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/Users/michham/github/AMO/.venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
NotImplementedError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__.py", line 27, in forward
    _4 = torch.slice(_3, 1, 0, 465)
    _5 = ops.prim.NumToTensor(torch.size(obs_teacher0, 0))
    _6 = torch.zeros([int(_5), 105], dtype=None, layout=None, device=torch.device("cuda:0"), pin_memory=False)
         ~~~~~~~~~~~ <--- HERE
    _7 = torch.slice(obs_teacher0, 0, 0, 9223372036854775807)
    _8 = torch.slice(_7, 1, 465, 9223372036854775807)

Traceback of TorchScript, original code (most recent call last):
save_jit.py(65): forward
/home/jialong/anaconda3/envs/loco/lib/python3.8/site-packages/torch/nn/modules/module.py(1522): _slow_forward
/home/jialong/anaconda3/envs/loco/lib/python3.8/site-packages/torch/nn/modules/module.py(1541): _call_impl
/home/jialong/anaconda3/envs/loco/lib/python3.8/site-packages/torch/nn/modules/module.py(1532): _wrapped_call_impl
/home/jialong/anaconda3/envs/loco/lib/python3.8/site-packages/torch/jit/_trace.py(1088): trace_module
/home/jialong/anaconda3/envs/loco/lib/python3.8/site-packages/torch/jit/_trace.py(820): trace
save_jit.py(215): save
save_jit.py(227): <module>
RuntimeError: Could not run 'aten::empty.memory_format' with arguments from the 'CUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'aten::empty.memory_format' is only available for these backends: [CPU, MPS, Meta, QuantizedCPU, QuantizedMeta, MkldnnCPU, SparseCPU, SparseMeta, SparseCsrCPU, SparseCsrMeta, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradHIP, AutogradXLA, AutogradMPS, AutogradIPU, AutogradXPU, AutogradHPU, AutogradVE, AutogradLazy, AutogradMTIA, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, AutogradMeta, AutogradNestedTensor, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, BatchedNestedTensor, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PreDispatch, PythonDispatcher].

What I have read until now, this cannot be fixed on the calling side of the model. Can this only be run in a CUDA environment?

Thanks in advance for any help,
michael

[Tian-Shu-Huang]

The current model is saved on cuda, so it is not available to run on cpu.

[Asthestarsfalll]

Add this to support cpu:

    for node in policy_jit.graph.nodes():
        if node.kind() == "prim::Constant":
            node: torch.Node
            if "value" in node.attributeNames():
                value = node.output().toIValue()
                if isinstance(value, torch.device):
                    node.removeAttribute('value')
                    node.s_("value", 'cpu')

[michaelHampel]

@Asthestarsfalll , thanks a lot - now the model is running, very cool!!!