Source code for torch_lattice.nn.modules.crop

from typing import Optional, Tuple

from torch import nn

from torch_lattice import SparseTensor
from torch_lattice.nn import functional as F

__all__ = ["SparseCrop"]




class SparseCrop(nn.Module):
    def __init__(
        self,
        coords_min: Optional[Tuple[int, ...]] = None,
        coords_max: Optional[Tuple[int, ...]] = None,
    ) -> None:
        super().__init__()
        self.coords_min = coords_min
        self.coords_max = coords_max



    def forward(self, input: SparseTensor) -> SparseTensor:
        return F.spcrop(input, self.coords_min, self.coords_max)