La entrada incorrecta de CNN

Estoy entrenando una red neuronal ingresando lotes de imágenes. Cada lote contiene 32 imágenes de tamaño 224x224x3. Para el entrenamiento, utilizo este código:

model.train()
loss_over_time = []

for epoch in range(n_epochs):  

    running_loss = 0.0

    for batch_i, data in enumerate(trainLoader):

        images = data['image']
        n_input_channels = 1
        key_pts = data['keypoints']

        key_pts = key_pts.view(key_pts.size(0), -1)

        key_pts = key_pts.type(torch.FloatTensor)
        images = images.type(torch.FloatTensor)

        output_pts = model(images)

        loss = criterion(output_pts, key_pts)

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

        running_loss += loss.item()
        if batch_i % 32 == 31:    # print every 32 batches
            print('Epoch: {}, Batch: {}, Avg. Loss: {}'.format(epoch + 1, batch_i+1, running_loss/32.0))
            loss_over_time.append(running_loss/32.0)
            running_loss = 0.0

print('Finished Training')
return loss_over_time

Pero me enfrento al siguiente error:

RuntimeError: Given groups=1, weight of size [128, 1, 3, 3], expected input[32, 224, 224, 3] to have 1 channels, but got 224 channels instead

He buscado en StackOverflow y otros sitios web. Parece que definí el número de canales de entrada como (https://stackoverflow.com/questions/65057069/runtimeerror-given-groups-1-weight-of-size-32-1-3-3-expected-input1-3):

n_input_channels = 1

Pero no puedo encontrar una forma de hacerlo.