helper_functions.py: Setting the device back to it's original

[AADITI-NIRFARAKE]

def plot_decision_boundary(model: torch.nn.Module, X: torch.Tensor, y: torch.Tensor):
"""Plots decision boundaries of model predicting on X in comparison to y.

Source - https://madewithml.com/courses/foundations/neural-networks/ (with modifications)
""" 
device = X.device

# Put everything to CPU (works better with NumPy + Matplotlib)
model.to("cpu")
X, y = X.to("cpu"), y.to("cpu")

plot code

set the device back to the original

model.to(device)
X, y = X.to(device), y.to(device)