Trying to silence a new torch.load warning

src/dartsort/__init__.py

@@ -11,6 +11,7 @@
 from .transform import WaveformPipeline
 from .util.analysis import DARTsortAnalysis
 from .util.data_util import DARTsortSorting
+from .util.drift_util import registered_geometry
 from .util.waveform_util import make_channel_index
 from .cluster import merge, postprocess, density
 from . import util

src/dartsort/transform/all_transformers.py

@@ -1,3 +1,4 @@
+import torch
 from .amplitudes import AmplitudeFeatures, AmplitudeVector, MaxAmplitude, Voltage
 from .enforce_decrease import EnforceDecrease
 from .localize import Localization, PointSourceLocalization
@@ -6,6 +7,7 @@
 from .temporal_pca import TemporalPCADenoiser, TemporalPCAFeaturizer, TemporalPCA
 from .transform_base import Waveform, Passthrough
 from .decollider import Decollider
+from .pipeline import WaveformPipeline
 
 all_transformers = [
     Waveform,
@@ -26,3 +28,7 @@
 ]
 
 transformers_by_class_name = {cls.__name__: cls for cls in all_transformers}
+
+# serialization
+others = [WaveformPipeline, set, torch.nn.ModuleList, slice, torch.nn.Sequential]
+torch.serialization.add_safe_globals(all_transformers + others)

src/dartsort/transform/pipeline.py

@@ -2,9 +2,6 @@
 """
 import torch
 
-from .all_transformers import transformers_by_class_name
-
-
 class WaveformPipeline(torch.nn.Module):
     def __init__(self, transformers):
         super().__init__()
@@ -15,6 +12,7 @@ def __init__(self, transformers):
     def from_class_names_and_kwargs(
         cls, geom, channel_index, class_names_and_kwargs
     ):
+        from .all_transformers import transformers_by_class_name
         return cls(
             [
                 transformers_by_class_name[name](

src/dartsort/util/data_util.py

@@ -194,7 +194,11 @@ def get_tpca(sorting):
     """Look for the TemporalPCAFeaturizer in the usual place."""
     base_dir = sorting.parent_h5_path.parent
     model_dir = base_dir / f"{sorting.parent_h5_path.stem}_models"
-    pipeline = torch.load(model_dir / "featurization_pipeline.pt")
+    pipeline = torch.load(
+        model_dir / "featurization_pipeline.pt",
+        weights_only=True,
+        map_location="cpu",
+    )
     tpca = pipeline.transformers[0]
     return tpca
 

src/dartsort/util/nn_util.py

@@ -130,3 +130,9 @@ def __init__(self, module):
     def forward(self, inputs):
         waveforms, masks = inputs
         return self.module(waveforms), masks
+
+
+# is this what they want us to do??
+torch.serialization.add_safe_globals(
+    [ResidualForm, WaveformOnlyResidualForm, ChannelwiseDropout, Cat, Permute, WaveformOnly, nn.Flatten, nn.Linear, nn.Conv1d, nn.ReLU, nn.BatchNorm1d]
+)