class TorchCat2T(nn.Module):
"""Concatenate two tensors in channel dimension"""
def __init__(self):
super().__init__()
self.cat_func = FloatFunctional()
def forward(self, x, y):
return self.cat_func.cat([x, y], dim=1)
class TorchCat(nn.Module):
"""Wrapper around torch.cat so that all ops can be found at build"""
def __init__(self):
super().__init__()
self.cat_func = FloatFunctional()
def forward(self, tensors: List[torch.Tensor], dim: int):
return self.cat_func.cat(tensors, dim)