Commit
·
6181390
1
Parent(s):
da70fe9
cr
Browse files- handler.py +5 -0
handler.py
CHANGED
|
@@ -63,4 +63,9 @@ class EndpointHandler():
|
|
| 63 |
indices.append(batch.nonzero(as_tuple=True)[0].tolist())
|
| 64 |
vecs.append(batch[indices[-1]].tolist())
|
| 65 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 66 |
return [indices, vecs]
|
|
|
|
| 63 |
indices.append(batch.nonzero(as_tuple=True)[0].tolist())
|
| 64 |
vecs.append(batch[indices[-1]].tolist())
|
| 65 |
|
| 66 |
+
# release all the GPU memory cache that can be freed
|
| 67 |
+
if torch.cuda.is_available():
|
| 68 |
+
torch.cuda.empty_cache()
|
| 69 |
+
|
| 70 |
+
|
| 71 |
return [indices, vecs]
|