Commit 52609ff8 authored by Benjamin Vandersmissen's avatar Benjamin Vandersmissen
Browse files

Change to calculating the pruning threshold, as torch doesn't like quantiles with large tensors

parent 00fc26a9
......@@ -26,7 +26,7 @@ def prune_by_magnitude(model: torch.nn.Module, percentage=0.2):
else:
weights = torch.cat((weights, reshapen_weight))
quantile = torch.quantile(weights, percentage)
quantile = np.quantile(weights.detach().cpu().numpy(), percentage)
for layer in model.modules():
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment