[master] add TorchApplyRecorderMixin

Author: vacancy

jactorch/nn/simple.py

@@ -8,9 +8,10 @@
 # This file is part of Jacinle.
 # Distributed under terms of the MIT license.
 
+import torch
 import torch.nn as nn
 
-__all__ = ['Identity']
+__all__ = ['Identity', 'TorchApplyRecorderMixin']
 
 
 class Identity(nn.Module):
@@ -19,3 +20,19 @@ def forward(self, *args):
             return args[0]
         return args
 
+
+class TorchApplyRecorderMixin(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._apply_recorder_indicator = nn.Parameter(
+            torch.tensor(0, dtype=torch.float32, device=torch.device('cpu'))
+        )
+        self._apply_recorder_indicator.requires_grad = False
+
+    @property
+    def dtype(self):
+        return self._apply_recorder_indicator.dtype
+
+    @property
+    def device(self):
+        return self._apply_recorder_indicator.device