@register_torch_op(torch_alias=["grid_sampler"], override=True)
def torch_grid_sample(context, node):
    inputs = mil_get_inputs(context, node, expected=5)
    res = mb.grid_sample(
        input=inputs[0], 
        grid=inputs[1], 
        mode=inputs[2], 
        padding_mode=inputs[3], 
        align_corners=inputs[4],
        name=node.name
    )
    context.add(res)