How to use Optimizer State Sharding with Sharpness-Aware Minimization?

I am trying to setup OSS with sharpness-aware minimization (SAM).
https://github.com/davda54/sam
https://github.com/SamsungLabs/ASAM

I followed this guide to setup OSS, but I have difficulty wrapping SAM with OSS.
Due to both OSS and SAM being optimizer wrappers and having their own step functions, I am not sure how to combine both and call those functions. My initial clue is to wrap like this: OSS(SAM(Adam)).
https://fairscale.readthedocs.io/en/stable/tutorials/oss.html

Is there a minimal working example showing how to do that?