-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
2bf6427
commit 3370ac2
Showing
5 changed files
with
41 additions
and
16 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
XLA_FLAGS="--xla_gpu_cuda_data_dir=/usr/lib/cuda" | ||
LD_LIBRARY_PATH="/usr/lib/cuda/lib64:" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,5 +1,6 @@ | ||
tensorflow==2.10.0 | ||
pytest~=7.4.0 | ||
numpy~=1.25.1 | ||
pillow~=10.0.0 | ||
einops~=0.6.1 | ||
tensorflow==2.12.1 | ||
pytest~=7.3.1 | ||
numpy~=1.24.3 | ||
pillow~=9.4.0 | ||
einops~=0.6.1 | ||
python-dotenv~=1.0.0 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,11 +1,26 @@ | ||
# standard imports | ||
import os | ||
import pytest | ||
|
||
# third-party imports | ||
import pytest # type: ignore | ||
import tensorflow as tf # type: ignore | ||
|
||
|
||
@pytest.fixture(autouse=True) | ||
def use_cpu_only(): | ||
"""Conv2D and tf.nn.Conv2D are broken on my CUDA/cuDNN version for some reason.""" | ||
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' | ||
tf.config.set_visible_devices([], 'GPU') | ||
yield | ||
@pytest.fixture(params=['CPU', 'GPU'], autouse=True) | ||
def device(request): | ||
"""Selects a runtime device for Tensorflow, so when parametrized with both CPU & GPU, tests will be run for both.""" | ||
device_type = request.param | ||
if device_type == 'CPU': | ||
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' | ||
tf.config.set_visible_devices([], 'GPU') | ||
else: | ||
os.environ.pop('CUDA_VISIBLE_DEVICES', None) | ||
gpus = tf.config.experimental.list_physical_devices('GPU') | ||
if gpus: | ||
try: | ||
for gpu in gpus: | ||
tf.config.experimental.set_memory_growth(gpu, True) | ||
except RuntimeError: | ||
print("You should really try out this other framework called PyTorch") | ||
|
||
yield device_type |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters