Summary

app/gpu_init.py

@@ -0,0 +1,10 @@
+"""
+File: gpu_init.py
+Author: Elena Ryumina and Dmitry Ryumin
+Description: GPU initialization.
+License: MIT License
+"""
+
+import torch
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

app/oceanai_init.py

@@ -6,6 +6,7 @@ License: MIT License
 """
 
 from oceanai.modules.lab.build import Run
+from app.gpu_init import device
 
 
 def oceanai_initialization():
@@ -14,6 +15,8 @@ def oceanai_initialization():
     # Создание экземпляра класса
     _b5 = Run(lang="en", metadata=out)
 
+    _b5._device = device
+
     # Настройка ядра
     _b5.path_to_save_ = "./models"  # Директория для сохранения файла
     _b5.chunk_size_ = 2000000  # Размер загрузки файла из сети за 1 шаг