Marthee commited on
Commit
1fc9143
·
verified ·
1 Parent(s): 06feec2

Update doors_fasterrcnn.py

Browse files
Files changed (1) hide show
  1. doors_fasterrcnn.py +2 -2
doors_fasterrcnn.py CHANGED
@@ -175,8 +175,8 @@ def main_run(pdf_fullpath, weights_path, pdf_name):
175
  model = get_model(num_classes)
176
  # Load the saved model's state dictionary with map_location to handle CPU
177
  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
178
- model.load_state_dict(torch.load(weights_path, map_location=device))
179
- # model.load_state_dict(torch.load(weights_path, map_location=device), strict=False)
180
  # Set the model to evaluation mode
181
  model.eval()
182
  # Move the model to the appropriate device
 
175
  model = get_model(num_classes)
176
  # Load the saved model's state dictionary with map_location to handle CPU
177
  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
178
+ # model.load_state_dict(torch.load(weights_path, map_location=device))
179
+ model.load_state_dict(torch.load(weights_path))
180
  # Set the model to evaluation mode
181
  model.eval()
182
  # Move the model to the appropriate device