Bike stops Suddenly after 2 to 3 Km

I have a Royal Enfield 350 Classic,15000km and 4 years old.
These days I am facing a problem ,it stops after a few kms. I was on highway and it stopped so I visited a mechanic he said there is water in your petrol tank ,he cleaned petrol tank as well as carburettor.After few days again same problem, now this time I visited authorize service centre of Royal Enfield , they are saying there is a lot of rust in petrol tank and rust in your petrol on off switch ,so they changed the switch and cleaned the petrol tank as well. .But after taking my bike from the showroom again same problem.I don’t now why it occur ? now this time showroom guys saying the there is some problem in your engine so it is cutting fuel supply
Can someone please suggest me what could be the reason and what should i do?