Update pdfextract_fun.py
Browse files- pdfextract_fun.py +2 -2
pdfextract_fun.py
CHANGED
@@ -32,8 +32,8 @@ cfg.merge_from_file("cascade_dit_base.yml")
|
|
32 |
cfg.MODEL.WEIGHTS = "publaynet_dit-b_cascade.pth"
|
33 |
|
34 |
# Step 3: set device
|
35 |
-
|
36 |
-
cfg.MODEL.DEVICE = "cuda"
|
37 |
|
38 |
# Step 4: define model
|
39 |
predictor = DefaultPredictor(cfg)
|
|
|
32 |
cfg.MODEL.WEIGHTS = "publaynet_dit-b_cascade.pth"
|
33 |
|
34 |
# Step 3: set device
|
35 |
+
cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
36 |
+
#cfg.MODEL.DEVICE = "cuda"
|
37 |
|
38 |
# Step 4: define model
|
39 |
predictor = DefaultPredictor(cfg)
|