Skip to content

Commit 888f34b

Browse files
authored
Merge pull request #937 from hotchpotch/fix_device_visualized
Enable Visualized_BGE Inference on CPU
2 parents e0b614c + c96caae commit 888f34b

1 file changed

Lines changed: 2 additions & 0 deletions

File tree

FlagEmbedding/visual/modeling.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,8 @@ def __init__(self,
9696
if torch.cuda.is_available():
9797
self.device = torch.device('cuda')
9898
self.to(self.device)
99+
else:
100+
self.device = torch.device('cpu')
99101

100102
def load_model(self, model_weight):
101103
self.load_state_dict(torch.load(model_weight, map_location='cpu'))

0 commit comments

Comments
 (0)