14 lines
320 B
Python
14 lines
320 B
Python
import torch
|
|
|
|
model = torch.load("/data6/zhangw/ZW_gra_project/yolov5_result_pretrained/G0_cub100_yolov5s6_raw_e300/weights/best.pt", map_location="cpu")["model"]
|
|
model.eval()
|
|
|
|
x = torch.randn(1, 3, 640, 640)
|
|
y = model(x)
|
|
|
|
if isinstance(y, (list, tuple)):
|
|
for o in y:
|
|
print(o.shape)
|
|
else:
|
|
print(y.shape)
|