37 lines
957 B
Python
37 lines
957 B
Python
import os
|
|
import sys
|
|
import unittest
|
|
import torch.nn as nn
|
|
|
|
|
|
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
|
sys.path.append(MODULE_ROOT)
|
|
|
|
from pytorch_wrapper.utils import set_requires_grad
|
|
|
|
|
|
class TestSetRequiresGrad(unittest.TestCase):
|
|
def setUp(self):
|
|
self.model = nn.Linear(2, 3, bias=True)
|
|
|
|
def test_1(self):
|
|
# Check that gradients are initially enabled
|
|
for param in self.model.parameters():
|
|
self.assertTrue(param.requires_grad)
|
|
|
|
# Turn off gradients
|
|
set_requires_grad(self.model, False)
|
|
for param in self.model.parameters():
|
|
self.assertFalse(param.requires_grad)
|
|
|
|
# Turn on gradients
|
|
set_requires_grad(self.model, True)
|
|
for param in self.model.parameters():
|
|
self.assertTrue(param.requires_grad)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|