Procházet zdrojové kódy

Fix caffe2 Linux script

Lutz Roeder před 7 roky
rodič
revize
5810251b2e
1 změnil soubory, kde provedl 2 přidání a 1 odebrání
  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))