Sa-m commited on
Commit
742b795
·
1 Parent(s): dc93acd

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -0
app.py CHANGED
@@ -6,6 +6,11 @@ from tensorflow.keras.models import load_model
6
  import tensorflow_addons as tfa
7
  import os
8
  from tensorflow.keras.layers import *
 
 
 
 
 
9
 
10
  labels= {'Subway': 0, 'Starbucks': 1,'McDonalds': 2,'Burger King': 3,'KFC': 4,'Other': 5}
11
  HEIGHT,WIDTH=224,224
 
6
  import tensorflow_addons as tfa
7
  import os
8
  from tensorflow.keras.layers import *
9
+ #import torch.cuda
10
+
11
+
12
+ devices = tf.config.experimental.list_physical_devices('GPU')
13
+ tf.config.experimental.set_virtual_device_configuration(devices[0],tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
14
 
15
  labels= {'Subway': 0, 'Starbucks': 1,'McDonalds': 2,'Burger King': 3,'KFC': 4,'Other': 5}
16
  HEIGHT,WIDTH=224,224