Sfoglia il codice sorgente

Fix caffe2 Linux script

Lutz Roeder 7 anni fa
parent
commit
5810251b2e
1 ha cambiato i file con 2 aggiunte e 1 eliminazioni
  1. 2 1
      tools/caffe2-script.py

+ 2 - 1
tools/caffe2-script.py

@@ -179,7 +179,8 @@ def metadata():
                 update_argument(schema, arg)
             for input_desc in op_schema.input_desc:
                 update_input(schema, input_desc)
-            if operator_name != 'Int8ConvRelu' and operator_name != 'Int8AveragePoolRelu' and operator_name != 'Int8AveragePool':
+            skip_operator_output_map = { 'Int8ConvRelu': True, 'Int8Conv': True, 'Int8AveragePoolRelu': True, 'Int8AveragePool': True, 'Int8MaxPool': True }
+            if not operator_name in skip_operator_output_map:
                 for output_desc in op_schema.output_desc:
                     update_output(operator_name, schema, output_desc)
             schema['support_level'] = get_support_level(os.path.dirname(op_schema.file))