removing torch.cuda.empty_cache() from TF function (#3267)

torch.cuda.empty_cache() was being called from a TF function (even when torch is unavailable)
not sure any replacement is needed if TF OOMs
This commit is contained in:
Nitish Shirish Keskar authored and GitHub committed 2020-03-19 23:25:30 +01:00
1 parent ecfd336318
commit 8becb73293
1 file changed
-1
-1
View File
@@ -532,7 +532,6 @@ def _compute_tensorflow(
except tf.errors.ResourceExhaustedError as e:
print("Doesn't fit on GPU.", e)
torch.cuda.empty_cache()
dictionary[model_name]["results"][batch_size][slice_size] = "N/A"
dictionary[model_name]["memory"][batch_size][slice_size] = "N/A"
return dictionary