davanstrien HF Staff commited on
Commit
29dc431
·
1 Parent(s): fdbc942

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -2
app.py CHANGED
@@ -11,12 +11,20 @@ chk_point = "davanstrien/autotrain-ia-useful-covers-3665397856"
11
 
12
  model = AutoModelForImageClassification.from_pretrained(chk_point)
13
 
14
- pipe = pipeline(
 
15
  "image-classification",
16
  chk_point,
17
  accelerator="bettertransformer",device=device
18
  )
19
-
 
 
 
 
 
 
 
20
 
21
  def make_label_folders():
22
  folders = model.config.label2id.keys()
 
11
 
12
  model = AutoModelForImageClassification.from_pretrained(chk_point)
13
 
14
+ try:
15
+ pipe = pipeline(
16
  "image-classification",
17
  chk_point,
18
  accelerator="bettertransformer",device=device
19
  )
20
+ except NotImplementedError:
21
+ from transformers import pipeline
22
+ pipe = pipeline(
23
+ "image-classification",
24
+ chk_point,
25
+ device=device
26
+ )
27
+
28
 
29
  def make_label_folders():
30
  folders = model.config.label2id.keys()