Skip to content

Commit d8cc331

Browse files
Add Dropout layer to supported for sparse inference in XNNPACK
PiperOrigin-RevId: 374317714
1 parent 9fff2f3 commit d8cc331

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

tensorflow_model_optimization/python/core/sparsity/keras/pruning_policy.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,8 +160,9 @@ def _check_layer_support(self, layer):
160160
References:
161161
- https://github.com/google/XNNPACK/blob/master/src/subgraph.c#L130
162162
"""
163-
if isinstance(layer, (layers.Add, layers.Multiply, layers.ZeroPadding2D,
164-
layers.ReLU, layers.LeakyReLU, layers.ELU)):
163+
if isinstance(layer,
164+
(layers.Add, layers.Multiply, layers.ZeroPadding2D,
165+
layers.ReLU, layers.LeakyReLU, layers.ELU, layers.Dropout)):
165166
return True
166167
elif isinstance(layer, layers.DepthwiseConv2D):
167168
# 3x3 stride-1 convolution (no dilation, padding 1 on each side).

0 commit comments

Comments
 (0)