Function get_motion_module

src/models/motion_module.py:33–40

(in_channels, motion_module_type: str, motion_module_kwargs: dict)

def get_motion_module(in_channels, motion_module_type: str, motion_module_kwargs: dict):
    if motion_module_type == "Vanilla":
        return VanillaTemporalModule(
            in_channels=in_channels,
            **motion_module_kwargs,
        )
    else:
        raise ValueError


class VanillaTemporalModule(nn.Module):