def test_ir_fusion_fused_scale_momentum_decay():
        clear_files()
        context.set_context(save_graphs=True)
    
        epoch_size = 1
        batch_size = 1
        num_classes = 3
    
        input_np = np.random.uniform(0.0, 1.0,
                size=[batch_size, 3, 2, 2]).astype(np.float16)
        label_np = np.ones([batch_size, num_classes]).astype(np.float32)
        net = Net(3, num_classes)
        loss = SoftmaxCrossEntropyWithLogits(sparse=False)
        opt = Momentum(learning_rate=0.01, momentum=0.9,
                params=filter(lambda x: x.requires_grad, net.get_parameters()),
                weight_decay=1.5, loss_scale=1.5)
        lsm = FixedLossScaleManager(loss_scale=1.5, drop_overflow_update=False)
        net = amp.build_train_network(net, opt, loss,
                level="O3", loss_scale_manager=lsm)
        net.set_train()
        for epoch in range(0, epoch_size):
            net(Tensor(input_np), Tensor(label_np))
    
        result = find_files('hwopt*momentum_scale_fusion*ir',
                'FusedWeightScaleApplyMomentum')
>       assert result == '2'
E       AssertionError: assert '0' == '2'
E         - 0
E         + 2

def test_ir_fusion_combine_momentum_weight():
        clear_files()
        context.set_context(save_graphs=True)
    
        epoch_size = 1
        batch_size = 1
        num_classes = 3
    
        input_np = np.random.uniform(0.0, 1.0,
                size=[batch_size, 3, 2, 2]).astype(np.float16)
        label_np = np.ones([batch_size, num_classes]).astype(np.float32)
        net = Net2(3, num_classes)
        loss = SoftmaxCrossEntropyWithLogits(sparse=False)
        conv_params = list(filter(lambda x: 'conv' in x.name,
            net.trainable_params()))
        no_conv_params = list(filter(lambda x: 'conv' not in x.name,
            net.trainable_params()))
        group_params = [{'params': conv_params, 'weight_decay': 0.3},
                        {'params': no_conv_params, 'lr': 0.04},
                        {'order_params': net.trainable_params()}]
        opt = Momentum(group_params, learning_rate=0.03, momentum=0.9,
                loss_scale=1.3, weight_decay=0.7)
        lsm = FixedLossScaleManager(loss_scale=1.3, drop_overflow_update=False)
        net = amp.build_train_network(net, opt, loss, level="O3",
                loss_scale_manager=lsm)
        net.set_train()
        for epoch in range(0, epoch_size):
            net(Tensor(input_np), Tensor(label_np))
    
        result = find_files('hwopt*combine_momentum*ir',
                'CombineMomentumWeight')
>       assert result == '2'
E       AssertionError: assert '0' == '2'