mamkkl commited on
Commit
1cf8fb9
·
verified ·
1 Parent(s): 549ebd5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -8
app.py CHANGED
@@ -101,10 +101,10 @@ def loadModel():
101
  device_map="auto",
102
  #quantization_config=quantization_config,
103
  max_memory={
104
- 0: "20GB", # GPU 0 with 20GB memory
105
- 1: "40GB", # GPU 0 with 20GB memory
106
- 2: "40GB", # GPU 0 with 20GB memory
107
- 3: "40GB", # GPU 0 with 20GB memory
108
  #"cpu": "5GB" # CPU with 100GB memory
109
  },
110
  )
@@ -117,10 +117,10 @@ def loadModel():
117
  torch_dtype=torch.float16,
118
  is_trainable=False,
119
  max_memory={
120
- 0: "20GB", # GPU 0 with 20GB memory
121
- 1: "8GB", # GPU 0 with 20GB memory
122
- 2: "8GB", # GPU 0 with 20GB memory
123
- 3: "40GB", # GPU 0 with 20GB memory
124
  #"cpu": "5GB" # CPU with 100GB memory
125
  },
126
  )
 
101
  device_map="auto",
102
  #quantization_config=quantization_config,
103
  max_memory={
104
+ 0: "30GB", # GPU 0 with 20GB memory
105
+ 1: "45GB", # GPU 0 with 20GB memory
106
+ 2: "45GB", # GPU 0 with 20GB memory
107
+ 3: "45GB", # GPU 0 with 20GB memory
108
  #"cpu": "5GB" # CPU with 100GB memory
109
  },
110
  )
 
117
  torch_dtype=torch.float16,
118
  is_trainable=False,
119
  max_memory={
120
+ 0: "30GB", # GPU 0 with 20GB memory
121
+ 1: "45GB", # GPU 0 with 20GB memory
122
+ 2: "45GB", # GPU 0 with 20GB memory
123
+ 3: "45GB", # GPU 0 with 20GB memory
124
  #"cpu": "5GB" # CPU with 100GB memory
125
  },
126
  )