From eaecb5fc5775fa548e49a9f5c74c27272a3a3693 Mon Sep 17 00:00:00 2001 From: Akshay Kashyap Date: Mon, 13 May 2024 18:42:15 -0400 Subject: [PATCH] style --- lerobot/common/policies/diffusion/modeling_diffusion.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/lerobot/common/policies/diffusion/modeling_diffusion.py b/lerobot/common/policies/diffusion/modeling_diffusion.py index d284c38c..6a38a537 100644 --- a/lerobot/common/policies/diffusion/modeling_diffusion.py +++ b/lerobot/common/policies/diffusion/modeling_diffusion.py @@ -335,10 +335,7 @@ class SpatialSoftmax(nn.Module): # we could use torch.linspace directly but that seems to behave slightly differently than numpy # and cause a small degradation in pc_success of pre-trained models. - pos_x, pos_y = np.meshgrid( - np.linspace(-1., 1., self._in_w), - np.linspace(-1., 1., self._in_h) - ) + pos_x, pos_y = np.meshgrid(np.linspace(-1.0, 1.0, self._in_w), np.linspace(-1.0, 1.0, self._in_h)) pos_x = torch.from_numpy(pos_x.reshape(self._in_h * self._in_w, 1)).float() pos_y = torch.from_numpy(pos_y.reshape(self._in_h * self._in_w, 1)).float() # register as buffer so it's moved to the correct device, etc.