From 3664ba4e60555e9dd4046735cd7c1418d297e1b2 Mon Sep 17 00:00:00 2001
From: Max Lapan <max.lapan@gmail.com>
Date: Sun, 24 Nov 2019 18:00:57 +0300
Subject: [PATCH] Fix error with DeterministicPolicy

More pytorch-native way would be to use `Module.register_buffer()` method. In that case, buffer won't be used in parameters(), but will be converted to CUDA and CPU with `to()` call transparently.
---
 model.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/model.py b/model.py
index cc548e0..fcf53e6 100644
--- a/model.py
+++ b/model.py
@@ -148,4 +148,5 @@ class DeterministicPolicy(nn.Module):
     def to(self, device):
         self.action_scale = self.action_scale.to(device)
         self.action_bias = self.action_bias.to(device)
-        return super(GaussianPolicy, self).to(device)
+        self.noise = self.noise.to(device)
+        return super(DeterministicPolicy, self).to(device)