Error:
/tmp/ipykernel_35/2835307632.py:5: FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. Please use `torch.amp.GradScaler('cuda', args...)` instead.
  scaler = GradScaler()
/tmp/ipykernel_35/2835307632.py:28: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
  with autocast():
/opt/conda/lib/python3.10/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):
torch.Size([128, 16, 2])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[63], line 1
----> 1 trained_model = train_model(model, train_loader, val_loader, num_epochs_, learning_rate_)

Cell In[58], line 29, in train_model(model, train_loader, val_loader, num_epochs, learning_rate)
     27 print(batch_x.shape)
     28 with autocast():
---> 29     output = model(batch_x, decoder_input)
     30     loss = criterion(output, batch_y)
     31 # Scale the loss and call backward

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File /opt/conda/lib/python3.10/site-packages/torch/nn/parallel/data_parallel.py:186, in DataParallel.forward(self, *inputs, **kwargs)
    184     return self.module(*inputs[0], **module_kwargs[0])
    185 replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
--> 186 outputs = self.parallel_apply(replicas, inputs, module_kwargs)
    187 return self.gather(outputs, self.output_device)

File /opt/conda/lib/python3.10/site-packages/torch/nn/parallel/data_parallel.py:201, in DataParallel.parallel_apply(self, replicas, inputs, kwargs)
    200 def parallel_apply(self, replicas: Sequence[T], inputs: Sequence[Any], kwargs: Any) -> List[Any]:
--> 201     return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])

File /opt/conda/lib/python3.10/site-packages/torch/nn/parallel/parallel_apply.py:108, in parallel_apply(modules, inputs, kwargs_tup, devices)
    106     output = results[i]
    107     if isinstance(output, ExceptionWrapper):
--> 108         output.reraise()
    109     outputs.append(output)
    110 return outputs

File /opt/conda/lib/python3.10/site-packages/torch/_utils.py:706, in ExceptionWrapper.reraise(self)
    702 except TypeError:
    703     # If the exception takes multiple arguments, don't try to
    704     # instantiate since we don't know how to
    705     raise RuntimeError(msg) from None
--> 706 raise exception

RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/parallel/parallel_apply.py", line 83, in _worker
    output = module(*input, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/tmp/ipykernel_35/2193255145.py", line 67, in forward
    memory = self.encoder(src)  # (batch_size, src_seq_length, hidden_size)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/tmp/ipykernel_35/2193255145.py", line 26, in forward
    x = self.network(x)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/container.py", line 219, in forward
    input = module(input)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1603, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 308, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 304, in _conv_forward
    return F.conv1d(input, weight, bias, self.stride,
RuntimeError: Given groups=1, weight of size [8, 1, 7], expected input[64, 2, 16] to have 1 channels, but got 2 channels instead

Code:
# Function to train the model with teacher forcing
def train_model(model, train_loader, val_loader, num_epochs, learning_rate):
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scaler = GradScaler()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    for epoch in range(num_epochs):
        model.train()
        train_loss = 0.0
        for batch_x, batch_y in train_loader:
            if batch_y.dim() == 1:
                batch_x = batch_x.to(device).unsqueeze(-1)
                batch_y = batch_y.to(device).unsqueeze(-1)
            elif batch_y.dim() > 1:
                batch_x = batch_x.to(device)
                batch_y = batch_y.to(device)
            else:
                print("The input doesn't fit the expected shape of a uni or multi variate Timeseries")
            # Prepare decoder input by shifting targets and adding initial input token (teacher forcing)
            decoder_input = torch.zeros_like(batch_y)
            decoder_input[:, 1:, :] = batch_y[:, :-1, :]
            decoder_input[:, 0, :] = batch_x[:, -1, :]  # Initial input to decoder is last input of encoder

            optimizer.zero_grad()
            print(batch_x.shape)
            with autocast():
                output = model(batch_x, decoder_input)
                loss = criterion(output, batch_y)
            # Scale the loss and call backward
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            train_loss += loss.item()
        train_loss /= len(train_loader)

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for batch_x, batch_y in val_loader:
                if batch_y.dim() == 1:
                    batch_x = batch_x.to(device).unsqueeze(-1)
                    batch_y = batch_y.to(device).unsqueeze(-1)
                elif batch_y.dim() > 1:
                    batch_x = batch_x.to(device)
                    batch_y = batch_y.to(device)
                else:
                    print("The input doesn't fit the expected shape of a uni or multi variate Timeseries")

                decoder_input = torch.zeros_like(batch_y)
                decoder_input[:, 1:, :] = batch_y[:, :-1, :]
                decoder_input[:, 0, :] = batch_x[:, -1, :]
                with autocast():
                    output = model(batch_x, decoder_input)
                loss = criterion(output, batch_y)
                val_loss += loss.item()
            val_loss /= len(val_loader)

        #print(f"Epoch [{epoch+1}/{num_epochs}], Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

    return model