ensure trained_model is a Path before calling .suffix (#471)

brainglobe · Nov 14, 2024 · efc0f67 · efc0f67
1 parent 95abb3b
commit efc0f67
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/cellfinder/core/classify/classify.py b/cellfinder/core/classify/classify.py
@@ -1,5 +1,6 @@
 import os
 from datetime import datetime
+from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import keras
@@ -68,7 +69,7 @@ def main(
         workers=workers,
     )
 
-    if trained_model and trained_model.suffix == ".h5":
+    if trained_model and Path(trained_model).suffix == ".h5":
         print(
             "Weights provided in place of the model, "
             "loading weights into default model."