Fixes lint issue in AMP (apache#15015)

larroy · May 21, 2019 · aa55e3d · aa55e3d
1 parent f680255
commit aa55e3d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/mxnet/contrib/amp/amp.py b/python/mxnet/contrib/amp/amp.py
@@ -264,7 +264,7 @@ def init(target_dtype='float16', target_precision_ops=None,
     global _amp_initialized
     global _loss_scaler
     if not _amp_initialized:
-        assert target_dtype == 'float16' or target_dtype == np.float16, \
+        assert target_dtype in ['float16', np.float16], \
                "AMP currently supports only float16 as a target_dtype"
         _amp_initialized = True
         logging.info("Using AMP")