import torch
import torch.nn as nn

# Example model definition (use your actual architecture)
class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(64, 1, kernel_size=3, padding=1)

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.conv2(x)
        return x

# Step 1: Create the model instance
model = UNet()

# Step 2: Load the state dictionary
state_dict = torch.load('leukemia_cells_unet.pt', map_location=torch.device('cpu'), weights_only=True)

# Step 3: Load the weights into the model instance
model.load_state_dict(state_dict)

# Step 4: Set the model to evaluation mode
model.eval()