From c535b157070b3bc6180ab8160aef1838b7d3ab25 Mon Sep 17 00:00:00 2001 From: andreistirb Date: Tue, 15 Jan 2019 12:20:52 +0200 Subject: [PATCH] change the get_screen method so that it does not fill the gpu memory --- intermediate_source/reinforcement_q_learning.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/intermediate_source/reinforcement_q_learning.py b/intermediate_source/reinforcement_q_learning.py index 86e353e15e1..9a533f1983e 100644 --- a/intermediate_source/reinforcement_q_learning.py +++ b/intermediate_source/reinforcement_q_learning.py @@ -229,6 +229,7 @@ def conv2d_size_out(size, kernel_size = 5, stride = 2): # Called with either one element to determine next action, or a batch # during optimization. Returns tensor([[left0exp,right0exp]...]). def forward(self, x): + x = x.to(device) x = F.relu(self.bn1(self.conv1(x))) x = F.relu(self.bn2(self.conv2(x))) x = F.relu(self.bn3(self.conv3(x))) @@ -278,7 +279,7 @@ def get_screen(): screen = np.ascontiguousarray(screen, dtype=np.float32) / 255 screen = torch.from_numpy(screen) # Resize, and add a batch dimension (BCHW) - return resize(screen).unsqueeze(0).to(device) + return resize(screen).unsqueeze(0) env.reset()