fix demo flow

PythonOT · Oct 30, 2024 · 7f59a68 · 7f59a68
1 parent 108f9c8
commit 7f59a68
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/examples/others/plot_GMM_flow.py b/examples/others/plot_GMM_flow.py
@@ -56,8 +56,10 @@
 
 def draw_cov(mu, C, color=None, label=None, nstd=1, alpha=0.5):
     def eigsorted(cov):
+        if torch.is_tensor(cov):
+            cov = cov.detach().numpy()
         vals, vecs = np.linalg.eigh(cov)
-        order = vals.argsort()[::-1]
+        order = vals.argsort()[::-1].copy()
         return vals[order], vecs[:, order]
 
     vals, vecs = eigsorted(C)