From 363e0d315fb770c6d024c5aa979e50c3cbe0a1a8 Mon Sep 17 00:00:00 2001 From: Jannik Sven Scharrenbach Date: Tue, 9 Nov 2021 10:15:55 +0100 Subject: [PATCH] fixed KeyError attention_radius[(VEHICLE, VEHICLE)] in encode_edge --- trajectron/model/online/online_mgcvae.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/trajectron/model/online/online_mgcvae.py b/trajectron/model/online/online_mgcvae.py index c614c37a..f524e953 100644 --- a/trajectron/model/online/online_mgcvae.py +++ b/trajectron/model/online/online_mgcvae.py @@ -121,8 +121,8 @@ def _get_other_node_type_from_edge(self, edge_type_str): def _get_edge_type_from_str(self, edge_type_str): n1_type_str, n2_type_str = edge_type_str.split('->') - return (NodeType(n1_type_str, self.env.node_type_list.index(n1_type_str) + 1), - NodeType(n2_type_str, self.env.node_type_list.index(n2_type_str) + 1)) + return (self.env.NodeType[self.env.NodeType.node_type_list.index(n1_type_str)], + self.env.NodeType[self.env.NodeType.node_type_list.index(n2_type_str)]) def remove_edge_model(self, edge_type): if self.hyperparams['edge_encoding']: