I can’t understand how to change this code to convert it from the classifier of the whole 3d grid to a classifier of every cell in a grid.
Initial dimensionality of every grid is 16 x 12 x 10 and 2 classes for every cell
class Grids_Conv3d(nn.Module):
    def __init__(self, out_1=32, out_2=64):
        super(Grids_Conv3d, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv3d(in_channels=1, out_channels=out_1, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool3d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv3d(in_channels=out_1, out_channels=out_2, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool3d(kernel_size=2, stride=2))
        self.drop_out = nn.Dropout()
        self.fc1 = nn.Linear(1536, 1920)
        self.fc2 = nn.Linear(1920, 2) #number of classes

def forward(self, x):
    out = self.layer1(x)
    out = self.layer2(out)
    out = out.view(out.size(0), -1)
    out = self.drop_out(out)
    out = self.fc1(out)
    out = self.fc2(out)
    return F.log_softmax(out, dim=1)

Could anyone help me with this?
Thank you!