import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin

# Let's create two column vectors containing `0`s and `1`s.
batch = {'a': torch.tensor([[0], [0], [1], [1]]), 'b': torch.tensor([[0], [1], [0], [1]])}

class AND(nn.Module, PyTorchModelHubMixin):
    def __init__(self):
        super().__init__()
        self.operation = "C = A AND B"

    def forward(self, x):
        a = x['a']
        b = x['b']
        inputs = torch.cat([a, b], axis=1)
        column_sum = torch.sum(inputs, dim=1, keepdim=True)
        output = (column_sum >= 2).long()
        return output

# Instantiate:
logical_and = AND.from_pretrained("sadhaklal/and")

# Forward pass:
output = logical_and(batch)
print(output)