import torch import torchvision device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') model=torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True) model.load_state_dict(torch.load("oms_rcnn_weights.pth")) #pthファイルは適当 model.to(device) model
