fix the get_input_mask function

judge the layer name rather than idx

fix the get_input_mask function
judge the layer name rather than idx
3056bb86 · lishuang · 7aae9b81 · 3056bb86
隐藏空白更改
内联并排

Showing with 1 addition and 1 deletion

utils/prune_utils.py utils/prune_utils.py +1 -1

未找到文件。
--- a/utils/prune_utils.py
+++ b/utils/prune_utils.py
@@ -182,7 +182,7 @@ def get_input_mask(module_defs, idx, CBLidx2mask):

    if idx == 0:
        return np.ones(3)
-    if idx == 1:
+    if module_defs[idx - 1]['type'] == 'focus':
        return np.ones(12)
    if module_defs[idx - 1]['type'] == 'convolutional':
        return CBLidx2mask[idx - 1]