When it comes to the train loop, first I define the following:
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

class Learner(pl.LightningModule):
    def __init__(self, model:nn.Module):
        super().__init__()
        self.model = model
    
    def forward(self, x):
        return self.model(x)
    
    def training_step(self, batch, batch_idx):
        x, y = batch      
        y_hat = self.model(x)   
        loss = nn.MSELoss()(y_hat, y)
        logs = {'train_loss': loss}
        return {'loss': loss, 'log': logs}   
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.model.parameters(), lr=0.005)

    def train_dataloader(self):
        return trainloader

Then I define the model in the following way:
f = nn.Sequential(
        nn.Conv1d(1,64,2),
        nn.ReLU(),
        nn.MaxPool1d(2),
        nn.Flatten(start_dim=1),
        nn.Linear(64,50),
        nn.ReLU(),
        nn.Linear(50, 1)
        )

model = NeuralDE(f, sensitivity='adjoint', solver='dopri5').to(device)

and finally:
learn = Learner(model)
trainer = pl.Trainer(min_epochs=200, max_epochs=300)
trainer.fit(learn)