[ Mixed ] fix apply using casted function

[platform/core/ml/nntrainer.git] / nntrainer / layers / input_layer.cpp

diff --git a/nntrainer/layers/input_layer.cpp b/nntrainer/layers/input_layer.cpp
index 5417e44..70a056b 100644 (file)
--- a/nntrainer/layers/input_layer.cpp
+++ b/nntrainer/layers/input_layer.cpp
@@ -67,7 +67,13 @@ void InputLayer::exportTo(Exporter &exporter,
 }
 
 void InputLayer::finalize(InitLayerContext &context) {
-  context.setOutputDimensions(context.getInputDimensions());
+
+  std::vector<TensorDim> output_dims = context.getInputDimensions();
+
+  for (auto d : output_dims)
+    d.setTensorType({context.getFormat(), context.getActivationDataType()});
+
+  context.setOutputDimensions(output_dims);
 }
 
 } /* namespace nntrainer */