ModuleNotFoundError when inferencing mobile benchmark

[junjun1023]

Before Asking

• [X] I have read the README carefully. 我已经仔细阅读了README上的操作指引。

• [x] I want to train my custom dataset, and I have read the tutorials for training your custom data carefully and organize my dataset correctly; (FYI: We recommand you to apply the config files of xx_finetune.py.) 我想训练自定义数据集，我已经仔细阅读了训练自定义数据的教程，以及按照正确的目录结构存放数据集。（FYI: 我们推荐使用xx_finetune.py等配置文件训练自定义数据集。）

• [X] I have pulled the latest code of main branch to run again and the problem still existed. 我已经拉取了主分支上最新的代码，重新运行之后，问题仍不能解决。

Search before asking

• [X] I have searched the YOLOv6 issues and found no similar questions.

Question

I downloaded yolov6lite_s.pt. While running python tools/infer.py --weights ./weights/yolov6lite_s.pt --source images_dir --img-size 320, it report ModuleNotFoundError: No module named 'yolov6.models.yolo_lite'.

Here is the full log

Save directory already existed
Loading checkpoint from ./weights/yolov6lite_s.pt
Traceback (most recent call last):
  File "tools/infer.py", line 120, in <module>
    main(args)
  File "tools/infer.py", line 115, in main
    run(**vars(args))
  File "/home/kasm-user/anaconda3/envs/onepp/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "tools/infer.py", line 107, in run
    inferer = Inferer(source, webcam, webcam_addr, weights, device, yaml, img_size, half)
  File "/dataset/YOLOv6/yolov6/core/inferer.py", line 33, in __init__
    self.model = DetectBackend(weights, device=self.device)
  File "/dataset/YOLOv6/yolov6/layers/common.py", line 573, in __init__
    model = load_checkpoint(weights, map_location=device)
  File "/dataset/YOLOv6/yolov6/utils/checkpoint.py", line 25, in load_checkpoint
    ckpt = torch.load(weights, map_location=map_location)  # load
  File "/home/kasm-user/anaconda3/envs/onepp/lib/python3.8/site-packages/torch/serialization.py", line 789, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/kasm-user/anaconda3/envs/onepp/lib/python3.8/site-packages/torch/serialization.py", line 1131, in _load
    result = unpickler.load()
  File "/home/kasm-user/anaconda3/envs/onepp/lib/python3.8/site-packages/torch/serialization.py", line 1124, in find_class
    return super().find_class(mod_name, name)
ModuleNotFoundError: No module named 'yolov6.models.yolo_lite'

Does anyone encounter the same problem?

Additional

My torch version is 1.13.0.

[Chilicyy]

Hi @junjun1023, please check if you git pull the latest code of main branch and run inference under the directory of YOLOv6 project with setting the pythonpath of this project.