Add save function

borglab · Nov 29, 2020 · ac42a90 · ac42a90
1 parent fbe71f0
commit ac42a90
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/Sources/BeeTracking/AppearanceRAE.swift b/Sources/BeeTracking/AppearanceRAE.swift
@@ -338,6 +338,12 @@ public struct PretrainedDenseRAE: AppearanceModelEncoder {
     )
   }
 
+  /// Save the weight to file
+  public func save(to path: String) {
+    let np = Python.import("numpy")
+    np.save(path, np.array(inner.numpyWeights, dtype: Python.object))
+  }
+
   @differentiable
   public func encode(_ imageBatch: Tensor<Double>) -> Tensor<Double> {
     inner.encode(imageBatch)