Added freeze and fixed imports (#4)

Added freeze and fixed imports

Author: spietras

kilroy_module_pytorch_py_sdk/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "kilroy-module-pytorch-py-sdk"
-version = "0.2.0"
+version = "0.2.1"
 description = "SDK for kilroy modules using PyTorch 🧰"
 readme = "README.md"
 authors = ["kilroy <[email protected]>"]

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/__init__.py

@@ -30,11 +30,14 @@
     EpsilonNucleusCategoricalSampler,
 )
 from kilroy_module_pytorch_py_sdk.utils import (
+    truncate_first_element,
+    truncate_last_element,
     pad,
     unpad,
     pack_padded,
     pack_list,
     unpack_to_padded,
     unpack_to_list,
     squash_packed,
+    freeze,
 )

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/utils.py

@@ -1,4 +1,5 @@
 import weakref
+from contextlib import contextmanager
 from typing import (
     Any,
     Generic,
@@ -11,7 +12,7 @@
 )
 
 import torch
-from torch import Tensor
+from torch import Tensor, nn
 from torch.nn.utils.rnn import (
     PackedSequence,
     pack_padded_sequence,
@@ -88,3 +89,17 @@ def squash_packed(x, fn):
     return PackedSequence(
         fn(x.data), x.batch_sizes, x.sorted_indices, x.unsorted_indices
     )
+
+
+@contextmanager
+def freeze(model: nn.Module) -> nn.Module:
+    original_state = {}
+
+    for name, param in model.named_parameters():
+        original_state[name] = param.requires_grad
+        param.requires_grad = False
+
+    yield model
+
+    for name, param in model.named_parameters():
+        param.requires_grad = original_state[name]

pyproject.toml

@@ -3,7 +3,7 @@
 
 [tool.poetry]
 name = "kilroy-module-pytorch-py-sdk"
-version = "0.2.0"
+version = "0.2.1"
 description = "SDK for kilroy modules using PyTorch 🧰"
 readme = "kilroy_module_pytorch_py_sdk/README.md"
 authors = ["kilroy <[email protected]>"]