minor fix

alibaba · Jul 15, 2022 · 151b89a · 151b89a
1 parent 3be659f
commit 151b89a
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/federatedscope/core/auxiliaries/model_builder.py b/federatedscope/core/auxiliaries/model_builder.py
@@ -33,7 +33,13 @@ def get_shape_from_data(data, model_config, backend='torch'):
         num_label = data['num_label'] if 'num_label' in data else None
         num_edge_features = data[
             'num_edge_features'] if model_config.type == 'mpnn' else None
-        return (data['train'].x.shape, num_label, num_edge_features)
+        if model_config.task.startswith('graph'):
+            # graph-level task
+            data_representative = next(iter(data['train']))
+            return (data_representative.x.shape, num_label, num_edge_features)
+        else:
+            # node/link-level task
+            return (data.x.shape, num_label, num_edge_features)
 
     if isinstance(data, dict):
         keys = list(data.keys())