From b1632e4463fa1c220df48b5b5c8bce1d18beac51 Mon Sep 17 00:00:00 2001
From: dv-fenix <divyamgoel3@gmail.com>
Date: Thu, 26 Oct 2023 11:24:31 +0530
Subject: [PATCH] Updated square root computation for sparse matrices

---
 cleanfid/fid.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cleanfid/fid.py b/cleanfid/fid.py
index 8b8170a..e946b6e 100644
--- a/cleanfid/fid.py
+++ b/cleanfid/fid.py
@@ -43,13 +43,13 @@ def frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
     diff = mu1 - mu2
 
     # Product might be almost singular
-    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    covmean = linalg.fractional_matrix_power(sigma1.dot(sigma2), 0.5)
     if not np.isfinite(covmean).all():
         msg = ('fid calculation produces singular product; '
                'adding %s to diagonal of cov estimates') % eps
         print(msg)
         offset = np.eye(sigma1.shape[0]) * eps
-        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+        covmean = linalg.fractional_matrix_power((sigma1 + offset).dot(sigma2 + offset), 0.5)
 
     # Numerical error might give slight imaginary component
     if np.iscomplexobj(covmean):