2-8 学習データをC++から使いたい

2-8_SSD_inference.py で推論ができるところまでは確認しました。学習データをC++から使いたいと思っています。学習データを PyTorch C++ API から利用するために、「.pth」ファイルを「.pt」にして保存するため、 2-8_SSD_inference.py の最後に次のようなコードを追加してみました。

example = torch.rand(1, 3, 300, 300).to("cpu")
traced_net = torch.jit.trace(net, example)
torch.save(traced_net, '/tmp/a.pt')

これを実行した結果、次のようなエラーが出ました。

/home/soramimi/pytorch_advanced/2_objectdetection/utils/ssd_model.py:702: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if scores.nelement() == 0:  # nelementで要素数の合計を求めるTraceback (most recent call last):
  File "/home/soramimi/pytorch_advanced/2_objectdetection/./2-8_SSD_inference.py", line 122, in <module>
    traced_net = torch.jit.trace(net, example)
  File "/usr/lib/python3.10/site-packages/torch/jit/_trace.py", line 750, in trace
    return trace_module(
  File "/usr/lib/python3.10/site-packages/torch/jit/_trace.py", line 967, in trace_module
    module._c._create_method_from_trace(
  File "/usr/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1118, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/soramimi/pytorch_advanced/2_objectdetection/utils/ssd_model.py", line 813, in forward
    return self.detect.apply(output[0], output[1], output[2])
  File "/home/soramimi/pytorch_advanced/2_objectdetection/utils/ssd_model.py", line 715, in forward
    ids, count = nm_suppression(
  File "/home/soramimi/pytorch_advanced/2_objectdetection/utils/ssd_model.py", line 546, in nm_suppression
    keep = scores.new(scores.size(0)).zero_().long()
TypeError: expected TensorOptions(dtype=float, device=cpu, layout=Strided, requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt)) (got TensorOptions(dtype=long int, device=cpu, layout=Strided, requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt)))

「dtype=float」と「dtype=long int」という、データ型の違いのため、traceできないということのようです。

私はC++はだいたいわかりますが、Pythonについてはあまり詳しくないため、上のコードが正しいのか、そもそも、学習データをC++から利用するためのデータを作成することが可能なのかよくわからず、ネットで調べて試行錯誤しながら進めています。解決策はあるでしょうか？

【実行環境】

Arch Linux x86_64
CUDA 11.8
Python 3.10.8
PyTorch 1.12.1

YutaroOgawa / pytorch_advanced

2-8 学習データをC++から使いたい #212