From 860e4aa42182acc9fa4eefb054bce063543a2257 Mon Sep 17 00:00:00 2001
From: Bin Lu <lubing02@gmail.com>
Date: Tue, 24 Aug 2021 16:20:21 +0800
Subject: [PATCH] Update cosmargin.py

---
 ppcls/arch/gears/cosmargin.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/ppcls/arch/gears/cosmargin.py b/ppcls/arch/gears/cosmargin.py
index 51db5508..378e102a 100644
--- a/ppcls/arch/gears/cosmargin.py
+++ b/ppcls/arch/gears/cosmargin.py
@@ -46,6 +46,9 @@ class CosMargin(paddle.nn.Layer):
         weight = paddle.divide(weight, weight_norm)
 
         cos = paddle.matmul(input, weight)
+        if not self.training or label is None:
+            return cos
+        
         cos_m = cos - self.margin
 
         one_hot = paddle.nn.functional.one_hot(label, self.class_num)
-- 
GitLab