We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
2 parents e0b614c + c96caae commit 888f34bCopy full SHA for 888f34b
1 file changed
FlagEmbedding/visual/modeling.py
@@ -96,6 +96,8 @@ def __init__(self,
96
if torch.cuda.is_available():
97
self.device = torch.device('cuda')
98
self.to(self.device)
99
+ else:
100
+ self.device = torch.device('cpu')
101
102
def load_model(self, model_weight):
103
self.load_state_dict(torch.load(model_weight, map_location='cpu'))
0 commit comments