[Mixed] Support MSELoss - Mixed Precision

enable mixed precision in MSE Loss **Self evaluation:** 1. Build test: [X]Passed [ ]Failed [ ]Skipped 2. Run test: [X]Passed [ ]Failed [ ]Skipped Signed-off-by: Donghak PARK <[email protected]>
nnstreamer · May 23, 2024 · 83ee96e · 83ee96e
1 parent 3fe9a1e
commit 83ee96e
Showing 1 changed file with 24 additions and 2 deletions.
diff --git a/Applications/Custom/mae_loss.cpp b/Applications/Custom/mae_loss.cpp
@@ -24,15 +24,37 @@ static constexpr size_t SINGLE_INOUT_IDX = 0;
 
 void MaeLossLayer::forwarding(nntrainer::RunLayerContext &context,
                               bool training) {
-  nntrainer::Tensor &predicted = context.getInput(SINGLE_INOUT_IDX);
+  nntrainer::Tensor empty_tensor;
+
+  nntrainer::Tensor &predicted =
+    context.getInput(SINGLE_INOUT_IDX).getDataType() ==
+        ml::train::TensorDim::DataType::FP32
+      ? context.getInput(SINGLE_INOUT_IDX)
+      : empty_tensor;
+
+  if (predicted.empty())
+    predicted = context.getInput(SINGLE_INOUT_IDX)
+                  .clone(ml::train::TensorDim::DataType::FP32);
+
   nntrainer::Tensor &output = context.getOutput(SINGLE_INOUT_IDX);
 
   if (!context.executeInPlace())
     output.fill(predicted);
 }
 
 void MaeLossLayer::calcDerivative(nntrainer::RunLayerContext &context) {
-  nntrainer::Tensor &predicted = context.getInput(SINGLE_INOUT_IDX);
+  nntrainer::Tensor empty_tensor;
+
+  nntrainer::Tensor &predicted =
+    context.getInput(SINGLE_INOUT_IDX).getDataType() ==
+        ml::train::TensorDim::DataType::FP32
+      ? context.getInput(SINGLE_INOUT_IDX)
+      : empty_tensor;
+
+  if (predicted.empty())
+    predicted = context.getInput(SINGLE_INOUT_IDX)
+                  .clone(ml::train::TensorDim::DataType::FP32);
+
   nntrainer::Tensor &label = context.getLabel(SINGLE_INOUT_IDX);
 
   nntrainer::Tensor &deriv = context.getOutgoingDerivative(SINGLE_INOUT_IDX);