diff options
| -rw-r--r-- | ethosu/vela/tflite_writer.py | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py index 5aa38a7..6d85567 100644 --- a/ethosu/vela/tflite_writer.py +++ b/ethosu/vela/tflite_writer.py @@ -103,10 +103,10 @@ class TFLiteSerialiser: self.align_nng_inputs_to_tflite(op) all_ops.append(op) if op.type.is_conv2d_op() or op.type.is_depthwise_conv2d_op() or op.type == Op.FullyConnected: - # Op is run on CPU, make sure original tensor are written back + # Op is run on CPU, make sure the original weight and bias tensors are written back # instead of the cloned/reshaped (see tflite_reader) for idx, inp in enumerate(op.inputs): - if inp is not None and inp.src_tensor is not None: + if inp != op.ifm and inp is not None and inp.src_tensor is not None: op.inputs[idx] = inp.src_tensor # list of tuple(Op, string, op.version); the custom code is only used for 3rd party custom operators |