From 586de5e1d69369d40a7c084030714b85789c821b Mon Sep 17 00:00:00 2001
From: JW <48402711+CorleoneJW@users.noreply.github.com>
Date: Fri, 31 May 2024 14:55:41 +0800
Subject: [PATCH] Update lpips.py

When I using the lpips.LPIPS(net='vgg'), I found that the output would be the shape like [[[[0.3333]]]], which lead a bad using experience and the result could not be used in the last loss combination operation. However, add squeeze to delete the "[]" automatically could solve the problem mentioned above.
---
 lpips/lpips.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lpips/lpips.py b/lpips/lpips.py
index 93595476..e0cea74d 100755
--- a/lpips/lpips.py
+++ b/lpips/lpips.py
@@ -139,9 +139,9 @@ def forward(self, in0, in1, retPerLayer=False, normalize=False):
             val += res[l]
         
         if(retPerLayer):
-            return (val, res)
+            return (val.squeeze(), res)
         else:
-            return val
+            return val.squeeze()
 
 
 class ScalingLayer(nn.Module):