Update pixelization.py

pixelization.py

@@ -75,12 +75,12 @@ class Model():
             self.G_A_net = define_G(3, 3, 64, "c2pGen", "instance", False, "normal", 0.02, [0])
             self.alias_net = define_G(3, 3, 64, "antialias", "instance", False, "normal", 0.02, [0])
 
-            G_A_state = torch.load("160_net_G_A.pth" if not os.environ['NET_MODEL'] else os.environ['NET_MODEL'], map_location=str(self.device))
+            G_A_state = torch.load("160_net_G_A.pth" if not os.environ['NET_MODEL'] else os.environ['NET_MODEL'], map_location=str(self.device), weights_only=True)
             for p in list(G_A_state.keys()):
                 G_A_state["module."+str(p)] = G_A_state.pop(p)
             self.G_A_net.load_state_dict(G_A_state)
 
-            alias_state = torch.load("alias_net.pth" if not os.environ['ALIAS_MODEL'] else os.environ['ALIAS_MODEL'], map_location=str(self.device))
+            alias_state = torch.load("alias_net.pth" if not os.environ['ALIAS_MODEL'] else os.environ['ALIAS_MODEL'], map_location=str(self.device), weights_only=True)
             for p in list(alias_state.keys()):
                 alias_state["module."+str(p)] = alias_state.pop(p)
             self.alias_net.load_state_dict(alias_state)