i-dont-hug-face
commited on
Commit
•
26c91ca
1
Parent(s):
af679a9
Update inference.py
Browse files- inference.py +1 -0
inference.py
CHANGED
@@ -6,6 +6,7 @@ import io
|
|
6 |
import json
|
7 |
|
8 |
def model_fn(model_dir, context=None):
|
|
|
9 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
10 |
processor = DonutProcessor.from_pretrained(model_dir)
|
11 |
model = VisionEncoderDecoderModel.from_pretrained(model_dir)
|
|
|
6 |
import json
|
7 |
|
8 |
def model_fn(model_dir, context=None):
|
9 |
+
print("this is model_dir: ", model_dir)
|
10 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
11 |
processor = DonutProcessor.from_pretrained(model_dir)
|
12 |
model = VisionEncoderDecoderModel.from_pretrained(model_dir)
|