pep8
This commit is contained in:
1 parent
a7ea07036e
commit
82b35492af
2 files changed
+2
No files matched your search
@@ -41,6 +41,7 @@ def serve(
|
|||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
assert torch.cuda.is_available(), "CUDA not available"
|
assert torch.cuda.is_available(), "CUDA not available"
|
||||||
typer.echo("CUDA seems to be working")
|
typer.echo("CUDA seems to be working")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -38,6 +38,7 @@ USE_CUDA = os.environ.get("USE_CUDA_DOCKER", "false")
|
|||||||
if USE_CUDA.lower() == "true":
|
if USE_CUDA.lower() == "true":
|
||||||
try:
|
try:
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
assert torch.cuda.is_available(), "CUDA not available"
|
assert torch.cuda.is_available(), "CUDA not available"
|
||||||
DEVICE_TYPE = "cuda"
|
DEVICE_TYPE = "cuda"
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
Reference in new issue
Block a user