16 lines
437 B
Python
16 lines
437 B
Python
from lightning.pytorch.cli import LightningCLI
|
|
|
|
# from callbacks import SaveImageCallback
|
|
from datamodule import ColorDataModule
|
|
from model import ColorTransformerModel
|
|
|
|
|
|
def cli_main():
|
|
cli = LightningCLI(ColorTransformerModel, ColorDataModule) # noqa: F841
|
|
# note: don't call fit!!
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli_main()
|
|
# note: it is good practice to implement the CLI in a function and call it in the main if block
|