diff --git a/main.py b/main.py index 2a920a42c042b137bdf02dc0919718c9da4070fc..a34997546ac005845f63a2ca2ccda3e4c2d8a483 100644 --- a/main.py +++ b/main.py @@ -14,7 +14,7 @@ from unet_medical.unet_model import UNetMedical from nets.deeplab_v3 import deeplab_v3 from dataset import GetDatasetGenerator from loss import SoftmaxCrossEntropyLoss -import utils.learning_rates as learning_rates +from utils import learning_rates context.set_context(mode=context.PYNATIVE_MODE, save_graphs=False, device_target='Ascend', device_id=7)