[issue-787] added device parameter to the vision checks context

[yromanyshyn]

resolves #787

[ItayGabbay]

@yromanyshyn can you add tests?

[noamzbr]

Two things:

1. Need to merge that with the recent refactor in vision base
2. Did we test that this work on gpu?

[yromanyshyn]

@noamzbr

Two things:
Need to merge that with the recent refactor in vision base
Did we test that this work on gpu?

1. of course, currently I am writing tests for base vision classes and would like to add them to the current PR, in the near time I will push that changes.
2. I tried to run it on the collab and checks worked/ but that probably is not a sufficient enough test, will think about how to do it better.