From 829c635fe5c14712866d081aae03b98d2510f0d3 Mon Sep 17 00:00:00 2001
From: wangguanzhong <jerrywgz@126.com>
Date: Tue, 24 May 2022 10:36:19 +0800
Subject: [PATCH] enhance shm utils (#6043)

---
 ppdet/data/shm_utils.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/ppdet/data/shm_utils.py b/ppdet/data/shm_utils.py
index 38d8ba66c..a929a809c 100644
--- 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) * \
-- 
GitLab