test_sanity_import.py 284 Bytes
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

try:
    import transformer_engine.pytorch
    te_imported = True
except:
    te_imported = False

assert te_imported, 'transformer_engine import failed'
print("OK")