Change dim to self.dim in segmentation_models_pytorch.base.modules.ArgMax (#236)

Author: Luan Pham

segmentation_models_pytorch/base/modules.py

@@ -70,7 +70,7 @@ def __init__(self, dim=None):
         self.dim = dim
 
     def forward(self, x):
-        return torch.argmax(x, dim=dim)
+        return torch.argmax(x, dim=self.dim)
 
 
 class Activation(nn.Module):