Please give us fully-reproducible code - the minimal example with saving and loading the model.
Here is working example:
model = UNet(in_channels=3,
                      out_channels=1,
                      n_class=1,
                      kernel_size=3,
                      padding=1,
                      stride=1)
torch.save(model.state_dict(), "./test.pt")


new_model = model = UNet(in_channels=3,
                      out_channels=1,
                      n_class=1,
                      kernel_size=3,
                      padding=1,
                      stride=1)
new_model.load_state_dict(torch.load("./test.pt"))

x = torch.rand((1, 3, 224, 224))
print(model(x).equal(new_model(x)))