У меня есть код объекта Pytorch Dataset и обучения нейронной сети на двух GPU. Однако при обучении GPU задействуются очень слабо, изучи код и найди в нем слабые места. Можешь использовать profiler При необходимости.
class UserItemDataset(Dataset):
    def __init__(self, df, items_meta):
        self.df = df.reset_index(drop=True)
        self.items_meta = items_meta



        # Инициализация кэша
        self.cache_item_id = None
        self.cache_source_id = None
        self.cache_duration = None
        self.cache_embeddings = None

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # Получение основной записи
        row = self.df.iloc[idx]
        item_id = row['item_id']

        # Проверка кэша для item_id
        if item_id == self.cache_item_id:
            source_id = self.cache_source_id
            duration = self.cache_duration
            embeddings = self.cache_embeddings
        else:
            # Обновление кэша
            try:
                item_meta = self.items_meta.loc[item_id]
                source_id = item_meta['source_id']
                duration = item_meta['duration']
                embeddings = item_meta[embedding_cols].values.astype(np.float32)
            except KeyError:
                source_id = 0
                duration = 0.0
                embeddings = np.zeros(len(embedding_cols), dtype=np.float32)

            # Обновление кэша
            self.cache_item_id = item_id
            self.cache_source_id = source_id
            self.cache_duration = duration
            self.cache_embeddings = embeddings




        # Составление признаков
        gender = row['gender']
        age = row['age']
        # Целевая переменная
        target = row['target']

        # Преобразование в тензоры
        features = {
            'gender': torch.tensor(gender, dtype=torch.uint8),
            'age': torch.tensor(age, dtype=torch.uint8),
            'source_id': torch.tensor(source_id, dtype=torch.long),
            'duration': torch.tensor(duration, dtype=torch.uint8),
            'embeddings': torch.from_numpy(embeddings)
        }

        return features, torch.tensor(target, dtype=torch.long)

class ReactionNet(nn.Module):
    def __init__(self, num_gender=2, num_source=19613, embedding_dim=32):
        super(ReactionNet, self).__init__()
       
        self.source_embedding = nn.Embedding(num_source, 4)  # Размер embedding для source_id

        # Числовые признаки
        self.fc_num = nn.Sequential(
            nn.Linear(1 + 1 + 1 + 32, 64),  #  + age + duration +gender + embeddings
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, 32),
            nn.ReLU()
        )

        # Категориальные признаки
        self.fc_cat = nn.Sequential(
            nn.Linear(4, 16),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(16, 8),
            nn.ReLU()
        )

        # Объединение
        self.fc_combined = nn.Sequential(
            nn.Linear(32 + 8, 16),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(16, 3)  # Три класса: -1, 0, 1
        )

    def forward(self, x):
        # Категориальные Embeddings
        #gender_embedded = self.gender_embedding(x['gender'])
        #print(gender_embedded)
        source_embedded = self.source_embedding(x['source_id'])
        # Числовые признаки и эмбеддинги
        num_features = torch.stack([
            x['age'],
            x['gender'],
            x['duration']
        ], dim=1)
        embeddings = x['embeddings']
        num_combined = torch.cat([num_features, embeddings], dim=1)
        num_out = self.fc_num(num_combined)

        # Категориальные Features
        cat_features = source_embedded #torch.cat([gender_embedded, source_embedded], dim=1)
        cat_out = self.fc_cat(cat_features)

        # Объединение
        combined = torch.cat([num_out, cat_out], dim=1)
        out = self.fc_combined(combined)

        return out

from torch.utils.data import Dataset, DataLoader, Subset
from sklearn.model_selection import train_test_split

total_size = len(train)
train_size = int(0.8 * total_size)
val_size = int(0.2*total_size)

train_df, val_df = train[:train_size], train[train_size:]

train_df = train_df.sort_values(by='item_id')
val_df = val_df.sort_values(by='item_id')
train_dataset, val_dataset = UserItemDataset(train_df, items_meta), UserItemDataset(val_df, items_meta)

del train
batch_size = 2048  # Настройте в зависимости от объема памяти
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2,pin_memory=True)

# Инициализация модели
model = ReactionNet()

# Использование GPU, если доступно
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.nn.DataParallel(model, device_ids = [0,1]).to(device)
# Определение функции потерь и оптимизатора
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=3e-5)

from tqdm import tqdm
num_epochs = 2  # Настройте количество эпох

for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for batch in tqdm(train_loader):
        features, targets = batch
        for key in features:
            features[key] = features[key].to(device)
        targets = targets.to(device)
        # Обнуление градиентов
        optimizer.zero_grad()

        # Прямой проход
        outputs = model(features)
        loss = criterion(outputs, targets)

        # Обратный проход и оптимизация
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * targets.size(0)

    avg_loss = total_loss / train_size

    # Оценка на валидационной выборке
    model.eval()
    all_targets = []
    all_preds = []
    with torch.no_grad():
        for batch in tqdm(val_loader):
            features, targets = batch
            for key in features:
                features[key] = features[key].to(device)
            targets = targets.to(device)

            outputs = model(features)
            probabilities = nn.functional.softmax(outputs, dim=1)

            all_targets.extend(targets.cpu().numpy())
            all_preds.extend(probabilities.cpu().numpy())

    # Вычисление ROC AUC
    from sklearn.preprocessing import label_binarize

    y_true = label_binarize(all_targets, classes=[0, 1, 2])  # Имеется 3 класса: 0,1,2
    y_pred = np.array(all_preds)
    roc_auc = roc_auc_score(y_true, y_pred, average='macro', multi_class='ovr')

    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {avg_loss:.4f}, ROC AUC: {roc_auc:.4f}')