Requesting suggestion for AssertionError

[Abhijeet241093]

Cuda_Version : 10.1

Installed : conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=10.1 -c pytorch

Error :
python tools/dist_train.py --cfg experiments/coco/higher_hrnet/w32_512_adam_lr1e-3.yaml
-- Process 0 terminated with the following error:
Traceback (most recent call last):
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 19, in _wrap
fn(i, *args)
File "/media/phuongnd/SSD_DATA/HigherHRNet-Human-Pose-Estimation/tools/dist_train.py", line 196, in main_worker
writer_dict['writer'].add_graph(model, (dump_input, ))
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/tensorboardX/writer.py", line 901, in add_graph
self._get_file_writer().add_graph(graph(model, input_to_model, verbose))
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/utils/tensorboard/_pytorch_graph.py", line 286, in graph
trace = torch.jit.trace(model, args)
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 733, in trace
return trace_module(
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 928, in trace_module
module = make_module(mod, _module_class, _compilation_unit)
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 560, in make_module
return _module_class(mod, _compilation_unit=_compilation_unit)
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 1039, in init
tmp_module._modules[name] = make_module(
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 560, in make_module
return _module_class(mod, _compilation_unit=_compilation_unit)
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 1039, in init
tmp_module._modules[name] = make_module(
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 560, in make_module
return _module_class(mod, _compilation_unit=_compilation_unit)
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 1039, in init
tmp_module._modules[name] = make_module(
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 560, in make_module
return _module_class(mod, _compilation_unit=_compilation_unit)
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 1039, in init
tmp_module._modules[name] = make_module(
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 560, in make_module
return _module_class(mod, _compilation_unit=_compilation_unit)
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 1039, in init
tmp_module._modules[name] = make_module(
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 560, in make_module
return _module_class(mod, _compilation_unit=_compilation_unit)
File "/home/phuongnd/projects/anaconda3/envs/hrnet/lib/python3.8/site-packages/torch/jit/_trace.py", line 988, in init
assert isinstance(orig, torch.nn.Module)
AssertionError

[Abhijeet241093]

Error solved with cuda 12.4