[Train] Add support for handling multiple batch data types for prepare_data_loader

[VishDev12]

What happened + What you expected to happen

When working with Ray Train, using the ray.train.torch.prepare_data_loader method with a dataset that returns a dictionary instead of a tuple from its __getitem__ method causes issues.

That's because in the _WrappedDataLoader class, the _move_to_device method only deals with a tuple:

return tuple(try_move_device(i) for i in item)

ray/python/ray/train/torch/train_loop_utils.py

Lines 539 to 551 in 6dd564d

	def _move_to_device(self, item):
	if item is None:
	return None
	def try_move_device(i):
	try:
	i = i.to(self.device, non_blocking=self._auto_transfer)
	except AttributeError:
	logger.debug(f"Item {i} cannot be moved to device " f"{self.device}.")
	return i
	with torch.cuda.stream(self._memcpy_stream):
	return tuple(try_move_device(i) for i in item)

Ideally, this should deal with a dictionary type batch and an arbitrary level of nesting just like Pytorch does.

Versions / Dependencies

Python: 3.8.10
Ray: Tested with 1.13.0. Code hasn't changed in master except for a move from torch.py -> train_loop_utils.py

Reproduction script

Using any dataset that returns a dictionary from its getitem method should demonstrate this issue.

Issue Severity

Medium: It is a significant difficulty but I can work around it.

[xwjiang2010]

@amogkam

[truelegion47]

Hi folks, I have a Draft PR for this issue. Is this what we are intending to do here for dict handling ? Please provide feedback when you get a chance.

@amogkam Thanks in advance.