Adapted the code for Python 3.11 + Toch 2+

[treeizard]

What does this PR do?

• Modified the code slightly to support Python 3.11+ and newer torch versions.

Changes Made:

Change 1: Updated the torch device selection from Variable to torch.device
e.g. self.device = torch.device("cuda" if args.cuda and torch.cuda.is_available() else "cpu")

Change 2: Resolved the error for variable retrieval
e.g. batch_class_loss.item()

Change 3: Optimized model.py
e.g. x = x.flatten(start_dim = 1)

Evidence of full pipeline running:

Additional Note:

• Did not check for support on torch.distributed
• Thank you so much for making this project.