enhance shm utils (#6043)

829c635f · wangguanzhong · GitHub · c2d51430 · 829c635f
隐藏空白更改
内联并排

Showing with 4 addition and 1 deletion

ppdet/data/shm_utils.py ppdet/data/shm_utils.py +4 -1

未找到文件。
--- a/ppdet/data/shm_utils.py
+++ b/ppdet/data/shm_utils.py
@@ -34,7 +34,10 @@ SHM_DEFAULT_MOUNT = '/dev/shm'
 def _parse_size_in_M(size_str):
-    num, unit = size_str[:-1], size_str[-1]
+    if size_str[-1] == 'B':
+        num, unit = size_str[:-2], size_str[-2]
+    else:
+        num, unit = size_str[:-1], size_str[-1]
    assert unit in SIZE_UNIT, \
            "unknown shm size unit {}".format(unit)
    return float(num) * \