How to use the einops.layers.torch.Reduce function in einops

def create_model(use_reduce=False):
            return Sequential(
                Conv2d(3, 6, kernel_size=5),
                Reduce('b c (h h2) (w w2) -> b c h w', 'max', h2=2, w2=2) if use_reduce else MaxPool2d(kernel_size=2),
                Conv2d(6, 16, kernel_size=5),
                Reduce('b c (h h2) (w w2) -> b (c h w)', 'max', h2=2, w2=2),
                Linear(16 * 5 * 5, 120),
                ReLU(),
                Linear(120, 84),
                ReLU(),
                Linear(84, 10),
            )

def create_model(use_reduce=False):
            return Sequential(
                Conv2d(3, 6, kernel_size=5),
                Reduce('b c (h h2) (w w2) -> b c h w', 'max', h2=2, w2=2) if use_reduce else MaxPool2d(kernel_size=2),
                Conv2d(6, 16, kernel_size=5),
                Reduce('b c (h h2) (w w2) -> b (c h w)', 'max', h2=2, w2=2),
                Linear(16 * 5 * 5, 120),
                ReLU(),
                Linear(120, 84),
                ReLU(),
                Linear(84, 10),
            )

def __init__(self, pool_spatial="mean", pool_temporal="mean"):
        super().__init__()

        self.model = r2plus1d_34_32_ig65m(num_classes=359, pretrained=True, progress=True)

        self.pool_spatial = Reduce("n c t h w -> n c t", reduction=pool_spatial)
        self.pool_temporal = Reduce("n c t -> n c", reduction=pool_temporal)

def __init__(self, pool_spatial="mean", pool_temporal="mean"):
        super().__init__()

        self.model = r2plus1d_34_32_ig65m(num_classes=359, pretrained=True, progress=True)

        self.pool_spatial = Reduce("n c t h w -> n c t", reduction=pool_spatial)
        self.pool_temporal = Reduce("n c t -> n c", reduction=pool_temporal)