fixed for CNN 2D pre-processing

SMLC-NYSBC · Aug 19, 2023 · 7bcf26e · 7bcf26e
1 parent f3a537c
commit 7bcf26e
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tardis_pytorch/utils/predictor.py b/tardis_pytorch/utils/predictor.py
@@ -305,6 +305,7 @@ def build_NN(self, NN: str):
                 )
             else:
                 self.cnn = Predictor(
+                    checkpoint=self.checkpoint[0],
                     network="fnet",
                     subtype="32",
                     model_type="cryo_mem",
@@ -315,6 +316,7 @@ def build_NN(self, NN: str):
 
                 # Build DIST network with loaded pre-trained weights
                 self.dist = Predictor(
+                    checkpoint=self.checkpoint[0],
                     network="dist",
                     subtype="triang",
                     model_type="s3dis",