Wrapping of external PyTorch models

[sfoucher]

Hello there,

Is there a way to wrap external models into Flash? for instance TorchGeo has semantic segmentation models:

from torchgeo.trainers import SemanticSegmentationTask
torchgeo_task= SemanticSegmentationTask(segmentation_model= 'unet',
                               loss= 'ce',
                               encoder_name= 'resnet18',
                               encoder_weights= 'imagenet',
                               in_channels=4,
                               num_classes= 7,
                               ignore_zeros=0
                               )

SemanticSegmentationTask is already based on LigthningModule , is it possible to provide some minimal guidelines on how to proceed?

Thank you!