Improve repr for DropPath module

2 years ago · 07d0c4ae96
parent e27c16b8a0
commit 07d0c4ae96
1 changed files with 3 additions and 0 deletions
--- a/timm/models/layers/drop.py
+++ b/timm/models/layers/drop.py
@ -164,3 +164,6 @@ class DropPath(nn.Module):

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+    def extra_repr(self):
+        return f'drop_prob={round(self.drop_prob,3):0.3f}'