Simple regression with integer output

Hi, I have been studying and doing some pytorch lately in our project and I came up with this problem. I am trying to create a model with number of users (integer) as output and with inputs of workplace area, day of the year and time of the day. As of now i am doing great with floats as output but I have no idea how to implement integer outputs. My model is written below.

import torch
from torch import nn


class TestModel(nn.Module):
    def __init__(
            self, input_size: int = 3,
            output_size: int = 1,
            hidden_size: int = 8,
            n_hidden_layers: int = 4
        ):
        super().__init__()

        module_list = [nn.Linear(input_size, hidden_size), nn.ReLU()]
        for _ in range(n_hidden_layers):
            module_list.extend([nn.Linear(hidden_size, hidden_size), nn.ReLU()])
        module_list.append(nn.Linear(hidden_size, output_size))

        self.sequential = nn.Sequential(*module_list)


    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.sequential(x)

Any help on the model modification and how to train the model is much appreciated. Thank you.

I don’t think this easily works with integers. Why not just train “pretending” your integers are floats for training, and then for inferencing round the predicted float values?

1 Like

Thank you for your response. Actually, that’s what I did when I presented the my output to my team. SInce Im new to neural networks and pytorch, I thought there is a better way to do this because I thought my code was just a hack. Thank you very much. You directed my to the right direction.