thanks for the reply, I am still stuck at the issue, and the loss is just overflowing.
self.highway_layers is a list of highway layers which are implemented as:
class HighwayLayer(nn.Module):
def __init__(self, input_size, bias=-1):
    super(HighwayLayer, self).__init__()
    self.plain_layer = nn.Linear(input_size, input_size)
    self.transform_layer = nn.Linear(input_size, input_size)
    self.transform_layer.bias.data.fill_(bias)
    
def forward(self, x): # Has to get a hidden state? No.
    plain_layer_output = nn.functional.relu(self.plain_layer(x)) # Wanted variable got tensor 
    transform_layer_output = nn.functional.softmax(self.transform_layer(x))
    transform_value = torch.mul(plain_layer_output, transform_layer_output)
    carry_value = torch.mul((1 - transform_layer_output), x)
    return torch.add(carry_value, carry_value) # This returns the same size as input.

Please take a look, surely I am doing something wrong: I created the layers, checked if dimensions match, addded a hidden variable for recurrence, so it looks fine by me…but it is not working… thanks.