Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
ColossalAI
Commits
6754f1b7
".github/git@developer.sourcefind.cn:OpenDAS/colossalai.git" did not exist on "5dd573c6b6dcb37f1505ae2306a63e2ea4d388d0"
Unverified
Commit
6754f1b7
authored
Jun 06, 2022
by
Ziyue Jiang
Committed by
GitHub
Jun 06, 2022
Browse files
fix module utils bug (#1066)
parent
a0064407
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
2 additions
and
2 deletions
+2
-2
colossalai/nn/module_utils.py
colossalai/nn/module_utils.py
+2
-2
No files found.
colossalai/nn/module_utils.py
View file @
6754f1b7
...
...
@@ -14,7 +14,7 @@ def register_colo_module(module_type: type, colo_module: ColoModule):
def
is_colo_module
(
module
:
torch
.
nn
.
Module
):
global
_COLOSSAL_MODULES
for
module_type
in
_COLOSSAL_MODULES
.
keys
():
if
isinstance
(
type
(
module
)
,
module_type
):
if
isinstance
(
module
,
module_type
):
return
True
return
False
...
...
@@ -23,7 +23,7 @@ def get_colo_module(module: torch.nn.Module):
global
_COLOSSAL_MODULES
if
is_colo_module
(
module
):
for
module_type
,
colo_module
in
_COLOSSAL_MODULES
.
items
():
if
isinstance
(
type
(
module
)
,
module_type
):
if
isinstance
(
module
,
module_type
):
return
colo_module
else
:
return
None
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment