edits

app.py

@@ -29,10 +29,18 @@ 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)
 net_sfm.load_state_dict(state['state_dict'])
+dim_red_params_dict = {}
+for name, param in net_sfm.named_parameters():
+    if 'dim_reduction' in name:
+        dim_red_params_dict[name] = param
+
 
 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'] = 
 net_imagenet = fire_network.init_network(**state['net_params']).to(device)
-net_imagenet.load_state_dict(state2['state_dict'])
+net_imagenet.load_state_dict(state2['state_dict']) #, strict=False)
 
 # ---------------------------------------
 transform = transforms.Compose([