edits

app.py

@@ -24,7 +24,7 @@ scales = [2.0, 1.414, 1.0, 0.707, 0.5, 0.353, 0.25]
 
 device = 'cpu'
 
-# Load net
+# Load nets
 state = torch.load('fire.pth', map_location='cpu')
 state['net_params']['pretrained'] = None # no need for imagenet pretrained model
 net_sfm = fire_network.init_network(**state['net_params']).to(device)
@@ -37,8 +37,7 @@ for name, param in net_sfm.named_parameters():
 
 state2 = torch.load('fire_imagenet.pth', map_location='cpu')
 state2['net_params'] = state['net_params']
-state2['state_dict'] += dim_red_params_dict
-# state2['net_params'] = 
+state2['state_dict'] = dict(state2['state_dict'], **dim_red_params_dict);
 net_imagenet = fire_network.init_network(**state['net_params']).to(device)
 net_imagenet.load_state_dict(state2['state_dict']) #, strict=False)
 
@@ -51,21 +50,16 @@ transform = transforms.Compose([
 # ---------------------------------------
 
 # class ImgDataset(data.Dataset):
-
 #     def __init__(self, images, imsize):
 #         self.images = images
 #         self.imsize = imsize
 #         self.transform = transforms.Compose([transforms.ToTensor(), \
 #             transforms.Normalize(**dict(zip(["mean", "std"], net.runtime['mean_std'])))])
-
-
 #     def __getitem__(self, index):
 #         img = self.images[index]
 #         img.thumbnail((self.imsize, self.imsize), Image.Resampling.LANCZOS)
 #         print('after imresize:', img.size)
 #         return  self.transform(img)
-
-
 #     def __len__(self):
 #         return len(self.images)