From: Manu Mathew <a0393608@ti.com>
Date: Mon, 14 Sep 2020 07:49:56 +0000 (+0530)
Subject: use AddBlock for addition in fpn
X-Git-Url: https://git.ti.com/gitweb?p=jacinto-ai%2Fpytorch-jacinto-ai-devkit.git;a=commitdiff_plain;h=23cace5ecbf40a60fff164596c7ddf6d46c42db8

use AddBlock for addition in fpn
---

diff --git a/modules/pytorch_jacinto_ai/xvision/models/pixel2pixel/fpnlite_pixel2pixel.py b/modules/pytorch_jacinto_ai/xvision/models/pixel2pixel/fpnlite_pixel2pixel.py
index 2933252..9427f78 100644
--- a/modules/pytorch_jacinto_ai/xvision/models/pixel2pixel/fpnlite_pixel2pixel.py
+++ b/modules/pytorch_jacinto_ai/xvision/models/pixel2pixel/fpnlite_pixel2pixel.py
@@ -70,6 +70,7 @@ class FPNLitePyramid(torch.nn.Module):
         self.smooth_convs = torch.nn.ModuleList()
         self.shortcuts = torch.nn.ModuleList()
         self.upsamples = torch.nn.ModuleList()
+        self.adds = torch.nn.ModuleList()
 
         shortcut0 = self.create_shortcut(current_channels, decoder_channels, activation) \
             if (current_channels != decoder_channels) else None
@@ -94,6 +95,7 @@ class FPNLitePyramid(torch.nn.Module):
             upsample = xnn.layers.UpsampleWith(decoder_channels, decoder_channels, upstride, interpolation_type,
                                                interpolation_mode)
             self.upsamples.append(upsample)
+            self.adds.append(xnn.layers.AddBlock())
         #
     #
 
@@ -121,7 +123,7 @@ class FPNLitePyramid(torch.nn.Module):
             x_s = shortcut(x_s)
             # updample current output and add to that
             x = upsample(x)
-            x = x + x_s
+            x = self.adds[idx]((x,x_s))
             # smooth conv
             y = smooth_conv(x) if (smooth_conv is not None) else x
             # use smooth output for next level in inloop_fpn