Trace and Lightning view not displayed for Lightning 2.X

Hi,

I have the problem that the trace and the lightning view are not displayed for profiling with Lightning2.X.

I adapted the resnet50_profiler_api.py.py to be a minimum working example:

import torch.profiler
import lightning.pytorch as pl
from lightning.pytorch.profilers import PyTorchProfiler
import torchvision.models as models
import torchvision.transforms as T
import torchvision
import torch.utils.data
import torch.optim
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch



class Model(pl.LightningModule):
   def __init__(self, model, criterion, optimizer) -> None:
       super().__init__()
       self.model = model
       self.criterion = criterion
       self.optimizer = optimizer

   def training_step(self, train_batch: dict, batch_idx: int) -> torch.Tensor:
       inputs, labels = train_batch
       outputs = self.model(inputs)
       return self.criterion(outputs, labels)

   def configure_optimizers(self) -> torch.optim:
       return self.optimizer


cudnn.benchmark = True

transform = T.Compose([T.Resize(256), T.CenterCrop(224), T.ToTensor()])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                       download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32,
                                         shuffle=True, num_workers=4)

model = models.resnet50(pretrained=True)
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

model = Model(model, criterion, optimizer)

trainer = pl.Trainer(
   num_sanity_val_steps=0,
   devices=1,
   accelerator="gpu",
   profiler=PyTorchProfiler(filename="profiling")

)
trainer.fit(model, train_dataloaders=trainloader)

print("done.")

requirements.txt:

pytorch-lightning==2.0.1.post0
tensorboard ==2.12.2
torchvision==0.15.1
torch-tb-profiler==0.4.1

Start by running the provided minimum working example with requirements.txt and inspect the output produced by PyTorchProfiler(filename="profiling"). Compare the generated profiler artifacts with the missing trace and Lightning view; done means both views are displayed when profiling with Lightning 2.X.