Set span for items in GridLayoutManager using SpanSizeLookup

I want to implement grid-like layout with section headers. Think of https://github.com/TonicArtos/StickyGridHeaders

What I do now:

mRecyclerView = (RecyclerView) view.findViewById(R.id.grid);
mLayoutManager = new GridLayoutManager(getActivity(), 2);
mLayoutManager.setSpanSizeLookup(new GridLayoutManager.SpanSizeLookup() {
			@Override
			public int getSpanSize(int position) {
				switch(mAdapter.getItemViewType(position)){
					case MyAdapter.TYPE_HEADER:
						return 1;
					case MyAdapter.TYPE_ITEM:
						return 2;
					default:
						return -1;
				}
			}
		});

mRecyclerView.setLayoutManager(mLayoutManager);

Now both regular items and headers have span size of 1. How do I solve this?

The problem was that header should have span size of 2, and regular item should have span size of 1. So correct implementations is:

mLayoutManager.setSpanSizeLookup(new GridLayoutManager.SpanSizeLookup() {
            @Override
            public int getSpanSize(int position) {
                switch(mAdapter.getItemViewType(position)){
                    case MyAdapter.TYPE_HEADER:
                        return 2;
                    case MyAdapter.TYPE_ITEM:
                        return 1;
                    default:
                        return -1;
                }
            }
        });

Header should have a span equal to the span count of the entire list.

mLayoutManager.setSpanSizeLookup(new GridLayoutManager.SpanSizeLookup() {
    @Override
    public int getSpanSize(int position) {
           switch(mAdapter.getItemViewType(position)){
                    case MyAdapter.TYPE_HEADER:
                        return mLayoutManager.getSpanCount();
                    case MyAdapter.TYPE_ITEM:
                        return 1;
                    default:
                        return -1;
                }
    }
});

Answer to my own question: Override the getSpanSizeLookup() from the Activity after setting the adapter.