cchallu / spectranet

Apache License 2.0
23 stars 7 forks source link

Data transfer between CPU and GPU #5

Open bcjanmay opened 1 year ago

bcjanmay commented 1 year ago

Hello,

At line 183 p_0_z = self.p_0_chains[idxs].to(idxs.device), this is an CPU objective, running this created errors.

We have replaced this line with a=self.p_0_chains.to(idxs.device)
p_0_z = a[idxs].to(idxs.device)

Now we are able to run the code on Google Colab GPU without any issues.

Please let us know if this change is the correct solution. Thanks.