I don’t think you are seeing a memory leak, but are allocating additional memory.
E.g. in your train_one_epoch you are using a pure FP32 forward pass and another one wrapped in autocast:
        # Iterate over train batches
        for step, data in enumerate(train_bar, start=1):

            # unpack
            inputs = data["X"].to(self.device, non_blocking=True)
            # .view(-1, 1) if BCELoss
            targets = data["y"].to(self.device, non_blocking=True).view(-1, 1)
            logits = self.model(inputs) # Forward pass logits

            batch_size = inputs.shape[0]
           
            if self.params.use_amp:
                self.optimizer.zero_grad()
                with torch.cuda.amp.autocast(enabled=True):
                    logits = self.model(inputs)  # Forward pass logits
                    #criterion = RMSELoss()
                    #curr_batch_train_loss = criterion(logits, targets)
                    curr_batch_train_loss = self.train_criterion(
                        targets,
                        logits,
                        CRITERION_PARAMS,
                    )

which would cause the additional memory usage as two computation graphs are computed.