From 64a880386a3f4866b49b85c9eb17fc4a765e0ae6 Mon Sep 17 00:00:00 2001 From: Yufei Zhu Date: Wed, 22 Feb 2023 16:22:32 +0100 Subject: [PATCH] fix when only one ped in scene --- trajnetbaselines/lstm/gridbased_pooling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trajnetbaselines/lstm/gridbased_pooling.py b/trajnetbaselines/lstm/gridbased_pooling.py index e1d89ed..30a1450 100644 --- a/trajnetbaselines/lstm/gridbased_pooling.py +++ b/trajnetbaselines/lstm/gridbased_pooling.py @@ -250,7 +250,7 @@ def occupancy(self, obs, other_values=None, past_obs=None): ## if only primary pedestrian present if num_tracks == 1: - return self.constant*torch.ones(1, self.pooling_dim, self.n, self.n, device=obs.device) + return self.constant*torch.ones(batch_size, self.pooling_dim, self.n, self.n, device=obs.device) ## Get relative position ## [batch_size, num_tracks, 2] --> [batch_size, num_tracks, num_tracks, 2]