def __init__(self, hidden_dim=256, input_dim=128, **kwargs):
    super(FlowHead, self).__init__(**kwargs)
    self.conv1 = create_update_Conv2d(
        c_in=input_dim, c_out=hidden_dim, k_size=3)
    self.conv2 = create_update_Conv2d(c_in=hidden_dim, c_out=2, k_size=3)