Refactoring docker images, updating CUDA to 12.6 and pytorch 2.6

This commit is contained in:
ekoc committed 2025-02-12 22:01:34 +01:00
1 parent 73d9d9450b
commit 284c100860
11 files changed
+123 -143

No files matched your search

+3
View File
@@ -2,6 +2,8 @@ import numpy as np
import re
import unicodedata
from fairseq import checkpoint_utils
from fairseq.data import Dictionary
import torch
import logging
@@ -66,6 +68,7 @@ def load_embedding(embedder_model):
#print("MODEL PATH IS", model_path)
# Load model ensemble and task
torch.serialization.add_safe_globals([Dictionary])
models = checkpoint_utils.load_model_ensemble_and_task(
[f"{model_path}"],
suffix="",