import megengine.module as M
import megengine.functional as F


class FlowHead(M.Module):
    def __init__(self, input_dim=128, hidden_dim=256):
        super(FlowHead, self).__init__()
        self.conv1 = 
M.Conv2d(input_dim, hidden_dim, 3, padding=1)