def __init__(self, args, hidden_dim=128, **kwargs):
    super(BasicUpdateBlock, self).__init__(**kwargs)

    self.args = args
    self.encoder = BasicMotionEncoder(args)
    self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128 + hidden_dim)
    self.flow_head = FlowHead(hidden_dim=256, input_dim=hidden_dim)
    if args.convex_upsampling:
      self.mask = tf.keras.Sequential(
          [create_update_Conv2d(c_in=128, c_out=256, k_size=3),
           tf.keras.layers.ReLU(),
           create_update_Conv2d(c_in=256, c_out=64 * 9, k_size=1)
          ])