text: """
  Original code from: https://github.com/NVIDIA/flownet2-pytorch
  Modified by: Siddharth Ancha
  """
  
  import torch.nn as nn
  import torch
  import numpy as np
  
  from flowmatch.networks.coord_conv import CoordConv2d, CoordConvTranspose2d
  
  
  def base_conv(batchNorm, in_planes, out_planes, kernel_size, stride, conv_layer):
      if batchNorm:
          return nn.Sequential(
              conv_layer(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=False),
              nn.BatchNorm2d(out_planes),
              nn.LeakyReLU(0.1,inplace=True)
          )
      else:
          return nn.Sequential(
              conv_layer(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=True),
              nn.LeakyReLU(0.1,inplace=True)
          )
  
  
  def conv(batchNorm, in_planes, out_planes, kernel_size=3, stride=1):
      return base_conv(batchNorm, in_planes, out_planes, kernel_size, stride, nn.Conv2d)
  
  
  def coord_conv(batchNorm, in_planes, out_planes, kernel_size=3, stride=1):
      return base_conv(batchNorm, in_planes, out_planes, kernel_size, stride, CoordConv2d)
  
  
  # def i_conv(batchNorm, in_planes, out_planes, kernel_size=3, stride=1, bias = True):
  #     if batchNorm:
  #         return nn.Sequential(
  #             nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=bias),
  #             nn.BatchNorm2d(out_planes),
  #         )
  #     else:
  #         return nn.Sequential(
  #             nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=bias),
  #         )
  
  
  def predict_flow(in_planes):
      return nn.Conv2d(in_planes,2,kernel_size=3,stride=1,padding=1,bias=True)
  
  
  def predict_flow_coord(in_planes):
      return CoordConv2d(in_planes, 2, kernel_size=3, stride=1, padding=1, bias=True)
  
  
  def base_deconv(in_planes, out_planes, deconv_layer):
      return nn.Sequential(
          deconv_layer(in_planes, out_planes, kernel_size=4, stride=2, padding=1, bias=True),
          nn.LeakyReLU(0.1,inplace=True)
      )
  
  
  def deconv(in_planes, out_planes):
      return base_deconv(in_planes, out_planes, nn.ConvTranspose2d)
  
  
  def coord_deconv(in_planes, out_planes):
      return base_deconv(in_planes, out_planes, CoordConvTranspose2d)
  
  
  class tofp16(nn.Module):
      def __init__(self):
          super(tofp16, self).__init__()
  
      def forward(self, input):
          return input.half()
  
  
  class tofp32(nn.Module):
      def __init__(self):
          super(tofp32, self).__init__()
  
      def forward(self, input):
          return input.float()
  
  
  def init_deconv_bilinear(weight):
      f_shape = weight.size()
      heigh, width = f_shape[-2], f_shape[-1]
      f = np.ceil(width/2.0)
      c = (2 * f - 1 - f % 2) / (2.0 * f)
      bilinear = np.zeros([heigh, width])
      for x in range(width):
          for y in range(heigh):
              value = (1 - abs(x / f - c)) * (1 - abs(y / f - c))
              bilinear[x, y] = value
      weight.data.fill_(0.)
      for i in range(f_shape[0]):
          for j in range(f_shape[1]):
              weight.data[i,j,:,:] = torch.from_numpy(bilinear)
  
  
  def save_grad(grads, name):
      def hook(grad):
          grads[name] = grad
      return hook
  
  '''
  def save_grad(grads, name):
      def hook(grad):
          grads[name] = grad
      return hook
  import torch
  from channelnorm_package.modules.channelnorm import ChannelNorm 
  model = ChannelNorm().cuda()
  grads = {}
  a = 100*torch.autograd.Variable(torch.randn((1,3,5,5)).cuda(), requires_grad=True)
  a.register_hook(save_grad(grads, 'a'))
  b = model(a)
  y = torch.mean(b)
  y.backward()
  '''