def unet_init(self):
'''
Initializes the UNet model and load the model weights
'''
model = UNet(1, scale_factor=1.0, resize_mode='bilinear')
model.to(self.device)
criterion = SoftDiceLoss()
criterion = self.to_device(criterion, self.device)
# load pretrained weights
model.load_state_dict(torch.load(self.unet_path, map_location=self.device))
self.model = model
def unet_init(self):
'''
Initializes the UNet model and load the model weights
'''
model = UNet(1, scale_factor=1.0, resize_mode='bilinear')
model.to(self.device)
criterion = SoftDiceLoss()
criterion = self.to_device(criterion, self.device)