My train function:-
def _train(self, cur_epoch, dataset_name):
        self.net.train()
        self.net.training = True
        self.scheduler = lr_scheduler.CosineAnnealingLR(
            self.optimizer, len(self.train_loader))  # cosine learning rate
        train_losses = 0.0
        clf_losses = 0.0
        metric_losses = 0.0
        d_losses = 0.0
        g_losses = 0.0
        correct = 0
        total = 0
        n_batch = len(self.train_loader)
        print(f'\n=> Training Epoch #{cur_epoch}')
        for batch_idx, (inputs, labels) in enumerate(self.train_loader):

            inputs, labels = inputs.to(
                self.device), labels.to(self.device)

            # seed_features = self.net.extract_features(inputs, seq_lens)
            features = inputs.float()

            labels = labels.type(torch.LongTensor)


            if self.hparams['manifold_mixup']:
                features, targets_a, targets_b, lam = mixup_data(features, labels,
                                                                 0.2, use_cuda=True)
                features, targets_a, targets_b = map(Variable, (features,
                                                                targets_a, targets_b))
            # apply pba transformation
            if self.hparams['use_modals']:
                # Generate features and the corresponding labels
                features,  new_labels = self.pm.apply_policy(features, labels, cur_epoch, batch_idx, verbose=1)
                features.to(self.device)

            outputs = self.net.forward(features)  # Forward Propagation

            # classification loss
            if self.hparams['mixup'] or self.hparams['manifold_mixup']:
                c_loss = mixup_criterion(
                    self.criterion, outputs, targets_a, targets_b, lam)
            else:
                c_loss = self.criterion(outputs, labels)  # Loss
            clf_losses += c_loss.item()

            # total loss
            loss = c_loss

            if self.hparams['metric_learning']:
                m_loss = self.metric_loss(features, labels)[0]
                metric_losses += m_loss.item()
                loss = self.metric_weight * m_loss + \
                       (1 - self.metric_weight) * c_loss

            train_losses += loss.item()

         

            self.optimizer.zero_grad()
            loss.backward()  # Backward Propagation
            clip_grad_norm_(self.net.parameters(), 5.0)
            self.optimizer.step()  # Optimizer update

           
            # Accuracy
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            if self.hparams['mixup']:
                correct += (lam * predicted.eq(targets_a.data).cpu().sum().float()
                            + (1 - lam) * predicted.eq(targets_b.data).cpu().sum().float())
            else:
                correct += (predicted == labels).sum().item()

        # step
        step = (cur_epoch - 1) * (len(self.train_loader)) + batch_idx
        total_steps = self.hparams['num_epochs'] * len(self.train_loader)

       

        return correct / total, train_losses / total

Lossess.py code
def forward(self, embeddings, target):
       triplets = self.triplet_selector.get_triplets(embeddings, target)
        
        if embeddings.is_cuda:
            triplets = triplets.cuda()
       
        ap_distances = (embeddings[triplets[:, 0]] -
                        embeddings[triplets[:, 1]]).pow(2).sum(1)  # .pow(.5)
        an_distances = (embeddings[triplets[:, 0]] -
                        embeddings[triplets[:, 2]]).pow(2).sum(1)  # .pow(.5)
        # print(f'+ve: {ap_distances.mean()}\t-ve: {an_distances.mean()}')
        losses = F.relu(ap_distances - an_distances + self.margin)
        # losses = torch.max(an_distances - ap_distances+ self.margin, 0)[0]

        return losses.mean(), len(triplets)

In the train function, I call losses.py at the line if self.param[‘metric learning’].
While executing following error occurs.
ap_distances = (embeddings[triplets[:, 0]] -
IndexError: too many indices for tensor of dimension 1
Any insights will be helpful!!!