Update main.py

README.md

@@ -1,6 +1,8 @@
 ## Federated Learning with Matched Averaging
 This is the code accompanying the ICLR 2020 paper "Federated Learning with Matched Averaging " Paper link: [https://openreview.net/forum?id=BkluqlSFDS]
 
+![Alt](https://repobeats.axiom.co/api/embed/705163eb0ad960618a2f6fedbb1947992ba5a212.svg "Repobeats analytics image")
+
 ### Overview
 ---
 FedMA algorithm is designed for federated learning of modern neural network architectures e.g. convolutional neural networks (CNNs) and LSTMs. FedMA constructs the shared global model in a layer-wise manner by matching and averaging hidden elements (i.e. channels for convolution layers; hidden states for LSTM; neurons for fully connected layers) with similar feature extraction signatures.

main.py

@@ -1288,7 +1288,7 @@ def BBP_MAP(nets_list, model_meta_data, layer_type, net_dataidx_map,
 
         type_of_this_layer = layer_type[2 * layer_index - 2]
         type_of_prev_layer = layer_type[2 * layer_index - 2 - 2]
-        first_fc_identifier = (('fc' in type_of_this_layer or 'classifier' in type_of_this_layer) and ('conv' in type_of_prev_layer or 'features' in type_of_this_layer))
+        first_fc_identifier = (('fc' in type_of_this_layer or 'classifier' in type_of_this_layer) and ('conv' in type_of_prev_layer or 'features' in type_of_prev_layer))
 
         if first_fc_identifier:
             first_fc_index = layer_index
@@ -1658,4 +1658,4 @@ def fedma_comm(batch_weights, model_meta_data, layer_type, net_dataidx_map,
                                  test_dl_global,
                                  assignments_list,
                                  comm_round=args.comm_round,
-                                 device=device)
+                                 device=device)