def test_tcn_forecaster_distributed(self):
    (train_data, val_data, test_data) = create_data()
    init_orca_context(cores=4, memory='2g')
    forecaster = LSTMForecaster(past_seq_len=24, input_feature_num=2, output_feature_num=2, loss='mae', lr=0.01, distributed=True)
    forecaster.fit(train_data, epochs=2)
    distributed_pred = forecaster.predict(test_data[0])
    distributed_eval = forecaster.evaluate(val_data)
    model = forecaster.get_model()
    assert isinstance(model, torch.nn.Module)
    forecaster.to_local()
    local_pred = forecaster.predict(test_data[0])
    local_eval = forecaster.evaluate(val_data)
    np.testing.assert_almost_equal(distributed_pred, local_pred, decimal=5)
    try:
        import onnx
        import onnxruntime
        local_pred_onnx = forecaster.predict_with_onnx(test_data[0])
        local_eval_onnx = forecaster.evaluate_with_onnx(val_data)
        np.testing.assert_almost_equal(distributed_pred, local_pred_onnx, decimal=5)
    except ImportError:
        pass
    model = forecaster.get_model()
    assert isinstance(model, torch.nn.Module)
    stop_orca_context()
def test_tcn_forecaster_distributed(self):
    (train_data, val_data, test_data) = create_data()
    init_orca_context(cores=4, memory='2g')
    forecaster = LSTMForecaster(past_seq_len=24, input_feature_num=2, output_feature_num=2, loss='mae', lr=0.01, distributed=True)
    forecaster.fit(train_data, epochs=2)
    distributed_pred = forecaster.predict(test_data[0])
    distributed_eval = forecaster.evaluate(val_data)
    model = forecaster.get_model()
    assert isinstance(model, torch.nn.Module)
    forecaster.to_local()
    local_pred = forecaster.predict(test_data[0])
    local_eval = forecaster.evaluate(val_data)
    np.testing.assert_almost_equal(distributed_pred, local_pred, decimal=5)
    try:
        import onnx
        import onnxruntime
        local_pred_onnx = forecaster.predict_with_onnx(test_data[0])
        local_eval_onnx = forecaster.evaluate_with_onnx(val_data)
        np.testing.assert_almost_equal(distributed_pred, local_pred_onnx, decimal=5)
    except ImportError:
        pass
    model = forecaster.get_model()
    assert isinstance(model, torch.nn.Module)
    stop_orca_context()