More general version check (awslabs#683)

vivicoco · Apr 18, 2019 · c2626ed · c2626ed
1 parent 69ed325
commit c2626ed
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sockeye/transformer.py b/sockeye/transformer.py
@@ -295,7 +295,7 @@ def get_valid_length_mask_for(data: mx.sym.Symbol,
     :param name: Name of symbol.
     :return: Bias symbol. Shape: (batch, seq_len)
     """
-    if mx.__version__ == "1.3.1":
+    if mx.__version__.startswith("1.3.1"):
         # TODO(fhieber): remove old branch eventually
         # mxnet 1.3.1's broadcast_like operator does not support individual axes yet. This branch uses another way
         # of creating the required zeros array.