modify grad clip

odlab/engine.py

@@ -71,11 +71,11 @@ def train_one_epoch(cfg,
 
         # Backward
         losses.backward()
-        if cfg.clip_max_norm > 0:
-            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.clip_max_norm)
 
         # Optimize
         if (iter_i + 1) % cfg.grad_accumulate == 0:
+            if cfg.clip_max_norm > 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.clip_max_norm)
             optimizer.step()
             optimizer.zero_grad()
 

odlab/models/basic/conv.py

@@ -61,7 +61,7 @@ class BasicConv(nn.Module):
             return self.act(self.norm(self.conv(x)))
         else:
             # Depthwise conv
-            x = self.norm1(self.conv1(x))
+            x = self.act(self.norm1(self.conv1(x)))
             # Pointwise conv
-            x = self.norm2(self.conv2(x))
+            x = self.act(self.norm2(self.conv2(x)))
             return x