def forward(self, input_states):
    input_states = input_states.to(dtype=torch.long)
    input_states = (input_states + 1)
    x_hidden = self.x_embedding(input_states[:, :, 0])
    y_hidden = self.y_embedding(input_states[:, :, 1])
    axis_hidden = torch.cat([x_hidden, y_hidden], dim=2)
    type_hidden = self.type_embedding(input_states[:, :, 2])
    if (self.pool_type == 'sum'):
        return (axis_hidden + type_hidden)
    elif (self.pool_type == 'con'):
        return torch.cat([axis_hidden, type_hidden], dim=2)
def forward(self, input_states):
    input_states = input_states.to(dtype=torch.long)
    input_states = (input_states + 1)
    x_hidden = self.x_embedding(input_states[:, :, 0])
    y_hidden = self.y_embedding(input_states[:, :, 1])
    axis_hidden = torch.cat([x_hidden, y_hidden], dim=2)
    type_hidden = self.type_embedding(input_states[:, :, 2])
    if (self.pool_type == 'sum'):
        return (axis_hidden + type_hidden)
    elif (self.pool_type == 'con'):
        return torch.cat([axis_hidden, type_hidden], dim=2)