Hello everybody,
I am trying to create a new training subset from two dataset, but I have a problem “Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same”.
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torchvision
from torchvision import datasets, models, transforms

import torch.nn.functional as F

import numpy as np
import pandas as pd #Agrego para trabajar matriz de resultados
import tqdm
import pickle # Utilizada para almacenar modelos

import os
import copy

import matplotlib.pyplot as plt
%matplotlib inline

class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)        

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1) # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
    
cnn = CNN() 

For example, the two dataset are:
##-----------------## CIFAR10
image_size = 32
batch_size = 4
classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Resize(image_size),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset_1 = torchvision.datasets.CIFAR10(root='./data', train=True,download=True, transform=transform)
testset_1 = torchvision.datasets.CIFAR10(root='./data', train=False,download=True, transform=transform)

trainloader_1 = torch.utils.data.DataLoader(trainset_1, batch_size=batch_size,shuffle=True, num_workers=2)
testloader_1 = torch.utils.data.DataLoader(testset_1, batch_size=batch_size,shuffle=False, num_workers=2)

##-----------------## FashionMNIST
image_size = 32
batch_size = 4
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot')

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Resize(image_size),
     transforms.Lambda(lambda x: x.repeat(3,1,1)),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset_2 = torchvision.datasets.FashionMNIST(root='./data/FashionMNIST', train=True, transform=transform, download=True)
testset_2 = torchvision.datasets.FashionMNIST(root='./data/FashionMNIST', train=False, download=True, transform=transform)

trainloader_2 = torch.utils.data.DataLoader(trainset_2, batch_size=batch_size, shuffle=True, num_workers=0) 
testloader_2 = torch.utils.data.DataLoader(testset_2, batch_size=batch_size, shuffle=False, num_workers=0)

Here is the problem, when I try to create the new set (now with 20 classes). Taking a subset of one and the other complete.
evens = list(range(0, len(trainset_1), 10))

sub_trainset_1 = torch.utils.data.Subset(trainset_1, evens)
sub_testset_1 = torch.utils.data.Subset(testset_1, evens)

trainset = torch.utils.data.ConcatDataset([sub_trainset_1, trainset_2]) 
testset = torch.utils.data.ConcatDataset([sub_testset_1, testset_2]) 

trainloader = torch.utils.data.DataLoader(trainset, shuffle=True, num_workers=0, batch_size=4)
testloader = torch.utils.data.DataLoader(testset, shuffle=False, num_workers=0, batch_size=4)

For train, I have:
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(cnn.parameters(), lr=0.001)

num_epochs = 5
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

for epoch in range(num_epochs):
  total_loss = 0
  total_examples = 0

  with tqdm.notebook.tqdm(total=len(trainloader), unit='batch', desc=f'Epoch {epoch+1}/{num_epochs}', 
                          position=100, leave=True) as pbar:  
    for X, Y in trainloader: 

      X = X.to(device)
      Y = Y.to(device)

      # Forward pass
      Y_ = cnn(X)
      loss = criterion(Y_, Y)

      # Backward pass
      optimizer.zero_grad()
      loss.backward()

      # Gradient descent
      optimizer.step()

      total_loss += loss.item()*X.size(0)
      total_examples += X.size(0)

      pbar.set_postfix(loss=total_loss/total_examples)
      pbar.update()

My output is:
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-29-2eebc2f5e475> in <module>
     14 
     15       # Forward pass
---> 16       Y_ = cnn(X)
     17       loss = criterion(Y_, Y)
     18 

c:\users\xyz\anaconda3\envs\envgene\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

<ipython-input-2-fbd28d5a9e42> in forward(self, x)
     10 
     11     def forward(self, x):
---> 12         x = self.pool(F.relu(self.conv1(x)))
     13         x = self.pool(F.relu(self.conv2(x)))
     14         x = torch.flatten(x, 1) # flatten all dimensions except batch

c:\users\xyz\anaconda3\envs\envgene\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

c:\users\xyz\anaconda3\envs\envgene\lib\site-packages\torch\nn\modules\conv.py in forward(self, input)
    397 
    398     def forward(self, input: Tensor) -> Tensor:
--> 399         return self._conv_forward(input, self.weight, self.bias)
    400 
    401 class Conv3d(_ConvNd):

c:\users\xyz\anaconda3\envs\envgene\lib\site-packages\torch\nn\modules\conv.py in _conv_forward(self, input, weight, bias)
    394                             _pair(0), self.dilation, self.groups)
    395         return F.conv2d(input, weight, bias, self.stride,
--> 396                         self.padding, self.dilation, self.groups)
    397 
    398     def forward(self, input: Tensor) -> Tensor:

RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

Do you know this problem? I have tried without good results.
Thankful for the feedback
K.