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.