milesial / Pytorch-UNet

PyTorch implementation of the U-Net for image semantic segmentation with high quality images
GNU General Public License v3.0
9.32k stars 2.52k forks source link

Run the model in CPU #508

Open akote123 opened 3 months ago

akote123 commented 3 months ago

I want to run the unet model in cpu. I have tried net = torch.hub.load('milesial/Pytorch-UNet', 'unet_carvana', pretrained=True, scale=0.5) But I am getting error as map this to cpu when I try to add mappings in torch.hub.load argument it says the argument is not there. So how can I run this model in CPU for inference