add tensor save/load in numpy format (#176)

microsoft · Aug 3, 2022 · bffee7c · bffee7c
1 parent f9ea8b3
commit bffee7c
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/tutel/system.py b/tutel/system.py
@@ -76,3 +76,14 @@ def record_time(is_cuda=None):
         import torch
         torch.cuda.synchronize()
     return time.time()
+
+def save(t, path):
+    import numpy as np
+    npv = t.detach().cpu().numpy()
+    np.save(path, npv)
+
+def load(path, device=None):
+    import numpy as np
+    import torch
+    npv = np.load(path)
+    return torch.tensor(npv, device=device)