no COCO model loaded

modules/train.py

@@ -37,10 +37,7 @@ def get_arrow_model(num_classes, num_keypoints=2):
     """
     # Load a model pre-trained on COCO, initialized without pre-trained weights
     device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
-    if device == torch.device('cuda'):
-        model = keypointrcnn_resnet50_fpn(weights=KeypointRCNN_ResNet50_FPN_Weights.COCO_V1)
-    else:
-        model = keypointrcnn_resnet50_fpn(weights=None)
+    model = keypointrcnn_resnet50_fpn(weights=None)
 
     # Get the number of input features for the classifier in the box predictor.
     in_features = model.roi_heads.box_predictor.cls_score.in_features
@@ -66,7 +63,7 @@ def get_faster_rcnn_model(num_classes):
     - model (torch.nn.Module): The modified Faster R-CNN model.
     """
     # Load a pre-trained Faster R-CNN model
-    model = fasterrcnn_resnet50_fpn(weights=FasterRCNN_ResNet50_FPN_Weights.COCO_V1)
+    model = fasterrcnn_resnet50_fpn(weights=None)
 
     # Get the number of input features for the classifier in the box predictor
     in_features = model.roi_heads.box_predictor.cls_score.in_features