From a37bf1bc4239f78498d2229a0fd852c1ea5c3a57 Mon Sep 17 00:00:00 2001 From: Jiarui Fang Date: Mon, 14 Mar 2022 21:39:48 +0800 Subject: [PATCH] [hotfix] rm test_tensor_detector.py (#413) --- tests/test_utils/test_tensor_detector.py | 41 ------------------------ 1 file changed, 41 deletions(-) delete mode 100644 tests/test_utils/test_tensor_detector.py diff --git a/tests/test_utils/test_tensor_detector.py b/tests/test_utils/test_tensor_detector.py deleted file mode 100644 index 06b1c1846..000000000 --- a/tests/test_utils/test_tensor_detector.py +++ /dev/null @@ -1,41 +0,0 @@ -#!/usr/bin/env python -# -*- encoding: utf-8 -*- - -import torch -import torch.nn as nn - -from colossalai.utils import TensorDetector - -class MLP(nn.Module): - def __init__(self): - super().__init__() - self.mlp = nn.Sequential(nn.Linear(64, 8), - nn.ReLU(), - nn.Linear(8, 32)) - - def forward(self, x): - return self.mlp(x) - -def test_tensor_detect(): - - data = torch.rand(64, requires_grad=True).cuda() - data.retain_grad() - model = MLP().cuda() - - detector = TensorDetector(log='test', include_cpu=False, module=model) - - detector.detect() - out = model(data) - - detector.detect() - loss = out.sum() - detector.detect() - loss.backward() - detector.detect() - model = MLP().cuda() - detector.detect() - detector.close() - torch.cuda.empty_cache() - -if __name__ == '__main__': - test_tensor_detect() \ No newline at end of file