How to load the best weights of yolov6 after training?

vision
1

import torch
from ultralytics import YOLO
from yolov6.models.yolo import build_model

device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)

model_path = ‘path_to_weights.pt’

model = build_model(cfg=‘/path/yolov6n.yaml’, num_classes=1, device=device, fuse_ab=False, distill_ns=False)

model.torch.load(model_path)

model.eval()
.
.
.

I am performing object detection and I have trained my dataset on yolov6. Now, I want to use the best weights learned during training to perform another task but I am facing difficulty in loading these weights. Can you help me find the correct way to do this?

2

Could you describe what kind of errors you are seeing while trying to load the weights?
A minimal example of saving and loading the parameters and buffers would be:

sd = model.state_dict()
torch.save(sd, "checkpoint.pt")
...
sd = torch.load("checkpoint.pt")
model = MyModel()
model.load_state_dict(sd)