TypeError: model must be a LightningModule, got UNet

[formerlya]

run python train.py

Total number of training data: 6000
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
Traceback (most recent call last):
File "/home/UVM-Net-main/Ultra/train.py", line 75, in
main()
File "/home/UVM-Net-main/Ultra/train.py", line 71, in main
trainer.fit(model=model, train_dataloaders=trainloader)
File "/usr/local/lib/python3.9/dist-packages/lightning/pytorch/trainer/trainer.py", line 538, in fit
model = _maybe_unwrap_optimized(model)
File "/usr/local/lib/python3.9/dist-packages/lightning/pytorch/utilities/compile.py", line 123, in _maybe_unwrap_optimized
raise TypeError(f"model must be a LightningModule, got {type(model).__qualname__}")
TypeError: model must be a LightningModule, got UNet

I only modified the import issue
from uvmb import UVMB->from net.uvmb import UVMB
from unet_part import *->from net.unet_part import *

[formerlya]

torch 1.13.1+cu117
torchvision 0.14.1+cu117
python 3.9

[lwtgithublwt]

Author

@formerlya 你好，我遇到了同样的问题，不知道是否解决？

[Vibe53]

Author

@formerlya 你好，我遇到了同样的问题，不知道是否解决？

我也是一样的问题，你解决了吗

[BeiBei0116]

change "model=UNet(n_channels=3)" to "model = UnetModel()" in train.py