fix: HuberLoss workaround added (#1239)

catalyst-team · Jun 16, 2021 · dfb8d32 · dfb8d32
1 parent f69c292
commit dfb8d32
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -33,6 +33,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Engine hotfix for `predict_loader` ([#1235](https://github.com/catalyst-team/catalyst/pull/1235))
 - ([#1230](https://github.com/catalyst-team/catalyst/pull/1230))
   - Hydra hotfix due to `1.1.0` version changes
+- `HuberLoss` name conflict for pytorch 1.9 hotfix ([#1239](https://github.com/catalyst-team/catalyst/pull/1239))
 
 
 ## [21.05] - YYYY-MM-DD

diff --git a/catalyst/contrib/nn/criterion/__init__.py b/catalyst/contrib/nn/criterion/__init__.py
@@ -1,5 +1,6 @@
 # flake8: noqa
 
+import torch
 from torch.nn.modules.loss import *
 
 from catalyst.contrib.nn.criterion.ce import (
@@ -22,7 +23,10 @@
     GradientPenaltyLoss,
     MeanOutputLoss,
 )
-from catalyst.contrib.nn.criterion.huber import HuberLoss
+
+if torch.__version__ < "1.9":
+    from catalyst.contrib.nn.criterion.huber import HuberLoss
+
 from catalyst.contrib.nn.criterion.iou import IoULoss
 from catalyst.contrib.nn.criterion.trevsky import TrevskyLoss, FocalTrevskyLoss
 from catalyst.contrib.nn.criterion.lovasz import (