From 81120309acff37307b2226fbda12277ca1662f93 Mon Sep 17 00:00:00 2001
From: Alby M <1473644+Linux-cpp-lisp@users.noreply.github.com>
Date: Wed, 21 Apr 2021 13:42:20 -0600
Subject: Add .to()

---
 torch_ema/ema.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

(limited to 'torch_ema')

diff --git a/torch_ema/ema.py b/torch_ema/ema.py
index b3487cf..2aa3004 100644
--- a/torch_ema/ema.py
+++ b/torch_ema/ema.py
@@ -165,6 +165,28 @@ class ExponentialMovingAverage:
             if param.requires_grad:
                 param.data.copy_(c_param.data)
 
+    def to(self, device=None, dtype=None) -> None:
+        r"""Move internal buffers of the ExponentialMovingAverage to `device`.
+
+        Args:
+            device: like `device` argument to `torch.Tensor.to`
+        """
+        # .to() on the tensors handles None correctly
+        self.shadow_params = [
+            p.to(device=device, dtype=dtype)
+            if p.is_floating_point()
+            else p.to(device=device)
+            for p in self.shadow_params
+        ]
+        if self.collected_params is not None:
+            self.collected_params = [
+                p.to(device=device, dtype=dtype)
+                if p.is_floating_point()
+                else p.to(device=device)
+                for p in self.collected_params
+            ]
+        return
+
     def state_dict(self) -> dict:
         r"""Returns the state of the ExponentialMovingAverage as a dict."""
         # Following PyTorch conventions, references to tensors are returned:
-- 
cgit v1.2.3