1d multi-class classification in pytorch

Question

I think I must be fundamentally misunderstanding something about PyTorch model construction, but I'm hoping someone here has an idea of how to get at what I'm trying to achieve.

My inputs are DNA sequences of characters A, C, T or G of length 104, which I one-hot encode. Thus, the inputs are of size 4x104.

The outputs are either 0 or 1 for three different categories, i.e. [0, 1, 1] could be one set of outputs. I am trying to predict these outputs.

I used a mishmash of tutorials to try to get at something close to what I am envisioning, and I am trying to use a 1D convolution of size 3 in order to hopefully capture some element of which codon is being coded for at each position.

This is the model I have defined:

class CNN(nn.Module):
    def __init__(self):
        super().__init__()

        # Define the convolutional layers
        self.conv1 = nn.Conv1d(in_channels=4, out_channels=64, kernel_size=3, stride=1, padding=1)
        # Define the dense layers.
        self.fc1 = nn.Linear(64, 3)

    def forward(self, x):
        print(x.size(0), x.size(1))
        
        x = self.conv1(x)
        print('after conv1', x.size(0), x.size(1))

        x = torch.relu(x)
        print('after relu', x.size(0), x.size(1))
        
        x = x.view(x.size(1), -1)
        print('after view', x.size(0), x.size(1))
        
        x = self.fc1(x)
        print('after fc1', x.size(0), x.size(1))
    
        return x

And this is how I'm trying to train it:

model       = CNN()
lossFn      = torch.nn.MSELoss()
optimizer   = torch.optim.Adam(model.parameters(), lr=0.001)
max_epochs  = 10
test_fraction = round(.7*len(data))
train_fraction = round(.3*len(data))

device = 'cpu'

test_data_gen = zip(data[0:test_fraction], target[0:test_fraction])
train_data_gen = zip(data[test_fraction:], target[test_fraction:])

# Train the model
# loop over our epochs
model.train()
for e in range(0, max_epochs):
    # set the model in training mode
    # initialize the total training and validation loss
    losses = 0
    # loop over the training set
    for x, y in train_data_gen:
        # perform a forward pass and calculate the training loss
        pred = model(x)
        
        y = torch.from_numpy(y)
            
        loss = lossFn(pred, y.float())
        
        # zero out the gradients, perform the backpropagation step,
        # and update the weights
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # add the loss to the total training loss so far 
        losses += loss.item()
  

I of course want to have my predictions be of size 3x1 to be comparable to the actual outputs, but currently this does not seem to be the case though I thought the linear function might have taken care of condensing the values from the convolutional layer into my desired size.

For the first iteration, the print statements are:

4 104
after conv1 64 104
after relu 64 104
after view 104 64
after fc1 104 3

I apologize in advance for probably some deep conceptual misunderstanding of what is going on under the hood, but sincerely appreciate any advice on what I'm doing wrong and how to tackle this problem!