from transformers import PreTrainedModel
from .configuration_fc import FCConfig
from torch.nn import Linear

class FCModel(PreTrainedModel):
    config_class = FCConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = Linear(in_features=10, out_features=config.num_nodes)
    
    def forward(self, tensor):
        # Use as forward similar to forward in torch
        return self.model.forward(tensor)