can you change the below code to find perplexity of the output and also use ray tuner to optimize the tuner on validation perplexity instead of the validation loss

class TransAmTrainable(tune.Trainable):
    def __init__(self, config):
        super(TransAmTrainable, self).__init__(config)
        self.model = TransAm(feature_size=config["feature_size"],
                             num_layers=config["num_layers"],
                             dropout=config["dropout"],
                             nhead=config["nhead"])
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = self.model.to(self.device)
        self.criterion = nn.MSELoss()
        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=config["lr"])
        self.scheduler = StepLR(self.optimizer, step_size=1, gamma=config["lr_decay"])
        self.input_window = 100 # number of input steps
        self.output_window = 1 # number of prediction steps, in this model its fixed to one
        self.dataset = Data(self.input_window, self.output_window)
        self.train_data, self.val_data = self.dataset.get_data()
        self.epochs = config["epochs"]
        self.batch_size = config["batch_size"]
        self.best_val_loss = float('inf')  # Track best validation loss
        self.patience = 5  # Set patience for early stopping
        self.epochs_no_improve = 0  # Counter for epochs without improvement


    def step(self):  
        train_loss = 0.0
        for epoch in range(self.epochs):  # Adjust number of epochs based on configuration
            epoch_loss = self.train_batch(epoch)
            train_loss += epoch_loss

        val_loss = self.validation_loss()
        # Update Weights & Biases with training and validation loss
        wandb.log({"train_loss": train_loss / self.epochs, "val_loss": val_loss})
        self.scheduler.step()
        
        # Early stopping logic
        if val_loss < self.best_val_loss:
            self.best_val_loss = val_loss
            self.epochs_no_improve = 0
        else:
            self.epochs_no_improve += 1
            if self.epochs_no_improve >= self.patience:
                print("Early stopping triggered!")
                return {"loss": val_loss}  # Terminate training

        return {"loss": val_loss}

    def train_batch(self, current_epoch):
        self.model.train()
        total_loss = 0.0
        start_time = time.time()
        
        for batch_idx in range(0, len(self.train_data) - 1, self.batch_size):
            data, targets = self.dataset.get_batch(self.train_data, batch_idx, self.batch_size)
            data, targets = data.to(self.device), targets.to(self.device)
            self.optimizer.zero_grad()
            output = self.model(data)
            loss = self.criterion(output, targets)
            cur_loss = loss.item() # loss step by step
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 0.7)
            self.optimizer.step()
            total_loss += loss.item()
            elapsed = time.time() - start_time
            # Calculate, log and return loss for entire step 
            # print(f"| batch {batch_idx:5d}/{len(self.train_data) // self.batch_size:5d}"
            #       f"| lr {self.scheduler.get_last_lr()[0]:02.6f}| "
            #       f"| loss {cur_loss:5.5f} | ppl {math.exp(cur_loss):8.2f}")
            
            wandb.log({"step": batch_idx,
                       "lr": self.scheduler.get_last_lr()[0],
                       "elapsed_time_per_step(seconds)": elapsed * 1000,
                       "loss": cur_loss,
                       "perplexity": math.exp(cur_loss)  # Perplexity
                       })
        
        # Calculate, log and return average loss for entire epoch
        avg_train_loss = total_loss / (len(self.train_data) // self.batch_size)
        wandb.log({"epoch": current_epoch, 
                   "avg_loss": avg_train_loss,
                   "elapsed_time_per_batch(seconds)": (time.time() - start_time) * 1000})
        # print(f"Epoch: {current_epoch}, Average Loss: {avg_train_loss:.5f}, ")
        return avg_train_loss
    
    def validation_loss(self):
        self.model.eval()
        total_loss = 0.0
        with torch.no_grad():
            for batch_idx in range(len(self.val_data)):
                data, target = self.val_data[batch_idx]
                data, target = data.to(self.device), target.to(self.device)
                output = self.model(data)
                total_loss += self.criterion(output, target).item()
            avg_loss = total_loss / len(self.val_data)  # Average loss
        
        return avg_loss  # Return the average validation loss